#!/usr/bin/env python3
"""
motor_test.py - Yahboom 4-ch driver over UART, motors on M3 + M4.

  python3 motor_test.py hand   # motors stay OFF. Turn wheels by hand, watch the counts.
  python3 motor_test.py spin   # WHEELS OFF THE GROUND. Low PWM steps, one motor at a time.

Ctrl-C always stops the motors.
"""
import sys, time, threading
import serial

PORT = "/dev/serial0"
BAUD = 115200
CHANNELS = {"M3": 2, "M4": 3}      # index in the 4-value commands
TICKS_PER_REV = 30 * 11 * 4         # 520 motor: 1:30 gearbox, 11 pulses/rev, x4 edges
SPIN_STEPS = [800, 1200, 1600, 2000, 2500, -1600]


class Driver:
    def __init__(self):
        self.ser = serial.Serial(PORT, BAUD, timeout=0.05)
        self.lock = threading.Lock()
        self.totals = None      # latest cumulative counts [M1, M2, M3, M4]
        self.msgs = 0
        self.other = []         # last few lines that weren't $MAll, for debugging
        self.alive = True
        self.error = None       # set if the serial port fails in the reader thread
        threading.Thread(target=self._rx, daemon=True).start()
        self.send("$pwm:0,0,0,0#")
        self.send("$mtype:1#")          # 1 = Yahboom 520 motor
        self.send("$upload:1,0,0#")     # stream total encoder counts as $MAll:a,b,c,d#

    def send(self, s):
        self.ser.write(s.encode())
        time.sleep(0.1)

    def pwm(self, ch, value):
        v = [0, 0, 0, 0]
        v[ch] = int(value)
        self.ser.write(f"$pwm:{v[0]},{v[1]},{v[2]},{v[3]}#".encode())

    def stop(self):
        for _ in range(3):
            self.ser.write(b"$pwm:0,0,0,0#")
            time.sleep(0.02)

    def get(self):
        with self.lock:
            return self.totals, self.msgs, list(self.other)

    def _rx(self):
        buf = b""
        while self.alive:
            try:
                buf += self.ser.read(self.ser.in_waiting or 1)
            except serial.SerialException as e:
                self.error = str(e)
                return
            while b"#" in buf:
                line, buf = buf.split(b"#", 1)
                i = line.rfind(b"$")
                if i < 0:
                    continue
                msg = line[i + 1:].decode(errors="ignore")
                with self.lock:
                    if msg.startswith("MAll:"):
                        try:
                            self.totals = [int(x) for x in msg[5:].split(",")]
                            self.msgs += 1
                        except ValueError:
                            pass
                    else:
                        self.other = (self.other + [msg])[-5:]

    def close(self):
        self.stop()
        self.alive = False
        self.ser.close()


def hold(d, ch, value, seconds):
    """Keep re-sending the PWM so it holds even if the board has a command timeout."""
    end = time.perf_counter() + seconds
    while time.perf_counter() < end:
        if d.error:
            sys.exit(f"\nSerial error: {d.error}")
        d.pwm(ch, value)
        time.sleep(0.1)


def hand(d):
    base, _, _ = d.get()
    print("Motors are OFF.")
    print("Turn ONE wheel by hand, one full turn FORWARD (as if the robot drives toward the tape).")
    print("Then the other wheel. Ctrl-C when done.\n")
    while True:
        if d.error:
            sys.exit(f"\nSerial error: {d.error}")
        t, n, _ = d.get()
        m3 = t[2] - base[2]
        m4 = t[3] - base[3]
        print(f"M3 {m3:+7d} ({m3 / TICKS_PER_REV:+5.2f} rev)    "
              f"M4 {m4:+7d} ({m4 / TICKS_PER_REV:+5.2f} rev)    [{n} msgs]   ", end="\r", flush=True)
        time.sleep(0.1)


def spin(d):
    input("Are the wheels OFF the ground? Press Enter to start, Ctrl-C to abort... ")
    for name, ch in CHANNELS.items():
        print(f"\n--- {name} ---")
        for p in SPIN_STEPS:
            hold(d, ch, p, 0.5)                     # let it spin up
            t0, n0, _ = d.get(); s0 = time.perf_counter()
            hold(d, ch, p, 1.0)                     # measure over 1 s
            t1, n1, other = d.get(); s1 = time.perf_counter()
            cps = (t1[ch] - t0[ch]) / (s1 - s0)
            rpm = cps / TICKS_PER_REV * 60
            print(f"pwm {p:+5d}  ->  {cps:+8.0f} counts/s   {rpm:+6.1f} wheel rpm   {n1 - n0:4d} msgs")
        d.stop()
        time.sleep(1.0)
    _, _, other = d.get()
    print("\nnon-encoder lines from the board:", other if other else "none")
    print("Done.")


def main():
    if len(sys.argv) != 2 or sys.argv[1] not in ("hand", "spin"):
        sys.exit(__doc__)
    d = Driver()
    try:
        time.sleep(1.0)
        t, n, other = d.get()
        if t is None:
            print("No $MAll messages from the driver board after 1 s.")
            print("Other lines received:", other if other else "nothing at all")
            print("Check: board powered, Pi TX -> board RX, Pi RX -> board TX, GND shared, serial enabled.")
            return
        print(f"Driver board is talking ({n} encoder messages so far).\n")
        if sys.argv[1] == "hand":
            hand(d)
        else:
            spin(d)
    except KeyboardInterrupt:
        pass
    finally:
        d.close()
        print("\nmotors stopped")


if __name__ == "__main__":
    main()