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
from ColorSensorDiagnostics import ColorSensorDiagnostics
battery = BatteryDiagnostics(HUB)
motor = MotorDiagnostics(HUB, Motor)
colorsensor = ColorSensorDiagnostics(HUB, ColorSensor)
CLEARCONFIRM = input("Clear the console before proceeding? Y/N (default: yes): ")
if(CLEARCONFIRM == "Y" or CLEARCONFIRM == "y" or CLEARCONFIRM == "yes" or CLEARCONFIRM == ""):
    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("\nWhich diagnostic do you want to perform?")
    print("Enter 'b' for battery diagnostics")
    print("Enter 'm' for motor diagnostics")
    print("Enter 'cs' for color sensor 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 these over a period of 3 seconds and provide average 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
    elif choice == "cs":
        print("---------------------COLOR SENSOR DIAGNOSTICS---------------------")
        break
    else:
        print("Invalid choice. Please enter 'b', 'm', or 'q'.")