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")
print("""
 ███████████  █████ █████ ██████   █████   █████████   ██████   ██████ █████   █████████   █████████ 
▒▒███▒▒▒▒▒███▒▒███ ▒▒███ ▒▒██████ ▒▒███   ███▒▒▒▒▒███ ▒▒██████ ██████ ▒▒███   ███▒▒▒▒▒███ ███▒▒▒▒▒███
 ▒███    ▒███ ▒▒███ ███   ▒███▒███ ▒███  ▒███    ▒███  ▒███▒█████▒███  ▒███  ███     ▒▒▒ ▒███    ▒▒▒ 
 ▒██████████   ▒▒█████    ▒███▒▒███▒███  ▒███████████  ▒███▒▒███ ▒███  ▒███ ▒███         ▒▒█████████ 
 ▒███▒▒▒▒▒▒     ▒▒███     ▒███ ▒▒██████  ▒███▒▒▒▒▒███  ▒███ ▒▒▒  ▒███  ▒███ ▒███          ▒▒▒▒▒▒▒▒███
 ▒███            ▒███     ▒███  ▒▒█████  ▒███    ▒███  ▒███      ▒███  ▒███ ▒▒███     ███ ███    ▒███
 █████           █████    █████  ▒▒█████ █████   █████ █████     █████ █████ ▒▒█████████ ▒▒█████████ 
▒▒▒▒▒           ▒▒▒▒▒    ▒▒▒▒▒    ▒▒▒▒▒ ▒▒▒▒▒   ▒▒▒▒▒ ▒▒▒▒▒     ▒▒▒▒▒ ▒▒▒▒▒   ▒▒▒▒▒▒▒▒▒   ▒▒▒▒▒▒▒▒▒  
                                                                                                     
The free and open source diagnostics tool for the LEGO® Education SPIKE™ Prime robots, designed for FIRST Lego League.
Developed by Team 65266, Lego Dynamics.                                                
          """
          )
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()
        print("Battery diagnostics completed.")

    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'.")