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
import umath

hub = PrimeHub()
class MotorDiagnostics:
  def __init__(self):
      self.testmotor = None
  def stdev(self, vals):
    data = vals
    if len(data) < 2:
        return 0
    # Calculate the mean
    mean = sum(data) / len(data)

    # Calculate the variance (sum of squared differences from the mean, divided by n-1 for sample standard deviation)
    variance = sum([(x - mean) ** 2 for x in data]) / float(len(data) - 1)

    # Calculate the standard deviation (square root of the variance)
    
    std_dev_manual = umath.sqrt(variance)
    

    return (std_dev_manual)
  def health_score(self, desired, avg_speed, stdev_speed, avg_load):
    # Speed accuracy: penalize % error
    accuracy = max(0, 100 - abs(avg_speed - desired) / desired * 100)

    # Stability: penalize deviation relative to desired
    stability = max(0, 100 - (stdev_speed / desired) * 100)
    max_speed, max_accel, max_torque = self.testmotor.control.limits()
    # Load: penalize load percentage directly
    load_pct = (avg_load / max_torque) * 100 if max_torque else 0
    load_score = max(0, 100 - load_pct)


    # Final score: average of the three
    return (accuracy + stability + load_score) / 3
  def initializeMotor(self):
    valid_ports = {"A", "B", "C", "D", "E", "F"}
    while True:
      motorinput = input(
        "This test will run your motor at 3 speeds: 180, 540, and 1000 degrees per second.\n"
        "Please make sure your motor is not under any load (for example, your hand) during the test.\n"
        "If you want to test the wheel's load, note that this will affect the load measurements.\n"
        "Enter the port for the motor you would like to test (A, B, C, D, E, or F): "
      ).strip().upper()

      if motorinput in valid_ports:
        self.testmotor = Motor(Port[motorinput])
        print(f"Motor initialized on port {motorinput}.")
        break  # exit the loop once a valid port is entered
      else:
        print("Invalid port. Please enter A, B, C, D, or E, or F.")

  def testSpeed(self, speed):
    self.testmotor.reset_angle(0)
    
    motorspeeds = []
    motorloads = []
    target_angle = speed * 3

    self.testmotor.run_angle(speed, target_angle, Stop.HOLD, False)
    while abs(self.testmotor.speed()) > 5:
        wait(20)
        motorspeeds.append(self.testmotor.speed())
        motorloads.append(self.testmotor.load())
    
    max_speed, max_accel, max_torque = self.testmotor.control.limits()
    print("Maximum motor speed (deg/s):", max_speed)

    print("Desired motor speed (degrees per second): ", str(speed))
    if motorspeeds:
      avg = sum(motorspeeds) / len(motorspeeds)
      print("Average motor speed (degrees per second):", avg)
      print("Motor speed deviation (this measures how much your motor deviates from the average speed):", str(self.stdev(motorspeeds)))
    else:
      print("No speed samples collected.")
      avg = 0
    if motorloads:
      avgload = sum(motorloads) / len(motorloads)
      print("Average motor load (mNm):", avgload)
      print("Motor load deviation (this measures how much your motor load deviates from the average load):", str(self.stdev(motorloads)))
    else:
      print("No load samples collected.")
      avgload = 0
   
    score = self.health_score(speed, avg, self.stdev(motorspeeds), avgload)
    print("Health score for this test:", score)
    return score
  def fullTest(self):
    self.initializeMotor()
    test180 = self.testSpeed(180)
    test540 = self.testSpeed(540)
    test1000 = self.testSpeed(1000)
    final = (test180 + test540 + test1000) / 3
    print("Final motor health score:", final)
    # print final health scores here as a combination of the 3 speeds, include speeds individual scores and total score