from pybricks.hubs import PrimeHub
from pybricks.pupdevices import Motor, ColorSensor, UltrasonicSensor, ForceSensor
from pybricks.parameters import Button, Color, Direction, Port, Side, Stop
from pybricks.robotics import DriveBase
from pybricks.tools import wait, StopWatch
hub = PrimeHub()
from BatteryDiagnostics import BatteryDiagnostics
from MotorDiagnostics import MotorDiagnostics
battery = BatteryDiagnostics()
motor = MotorDiagnostics()
clearConfirmation = input("Do you want to clear the console before proceeding? Y/N (default: yes): ")
if(clearConfirmation == "Y" or clearConfirmation == "y" or clearConfirmation == "yes" or clearConfirmation == ""):
    print("Clearing console... \n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n")

while True:
    print("\nWhat diagnostic do you want to perform?")
    print("Enter 'b' for Battery diagnostics")
    print("Enter 'm' for Motor diagnostics")
    print("Enter 'q' to Quit")

    choice = input("Your choice: ").strip().lower()

    if choice == "b":
        print("-----------------------BATTERY DIAGNOSTICS-----------------------")
        print("This test will check the battery voltage and current. It will measure the voltage and current over a period of 3 seconds and provide average values and deviation values. Your voltage should be above 7800 mV for optimal performance.")
        input("Press Enter to begin the battery diagnostics.")
        battery.printAll()

    elif choice == "m":
        print("------------------------MOTOR DIAGNOSTICS------------------------")
        motor.fullTest()
        print("Motor diagnostics completed.")

    elif choice == "q":
        print("Diagnostics completed successfully. Exiting with code 0. Good luck in the robot game!")
        break

    else:
        print("Invalid choice. Please enter 'b', 'm', or 'q'.")