"""Write a single goal to motors 1 and 2 (double servo joint test).

Both motors receive the same goal value since they drive the same joint.
No MuJoCo model needed — just two servos clamped together.

Usage:
  python write_double_servo_test.py <goal_rad>
  python write_double_servo_test.py 0.5
  python write_double_servo_test.py -- -0.3   # negative values need --
"""
import math
import sys
import time

from scservo_sdk import PortHandler
import scservo_patch
from scservo_sdk.sms_sts import sms_sts

USB_PORT = "/dev/tty.usbserial-0001"
BAUD_RATE = 115200

MOTOR_IDS = [1, 2]
SIGNS = [1, 1]

TICKS_PER_REV = 4096
CENTER_TICK = 2048
TICKS_TO_RAD = (2.0 * math.pi) / TICKS_PER_REV
RAD_TO_TICKS = TICKS_PER_REV / (2.0 * math.pi)

SAFE_TICK_LO = CENTER_TICK - 1251  # 797
SAFE_TICK_HI = CENTER_TICK + 1251  # 3299

DEFAULT_SPEED = 360
DEFAULT_ACCEL = 120

ADDR_TORQUE_ENABLE = 40
ADDR_TEMP = 63
TEMP_WARN_C = 55
TEMP_ABORT_C = 65
TEMP_POLL_EVERY_N = 10

MONITOR_HZ = 20
MONITOR_TIMEOUT_S = 120.0


def rad_to_ticks(rad, sign):
    return int(round(CENTER_TICK + sign * rad * RAD_TO_TICKS))


def emergency_stop(sts):
    for mid in MOTOR_IDS:
        try:
            sts.writeTxRx(mid, ADDR_TORQUE_ENABLE, 1, [0])
        except Exception:
            pass


def main():
    if len(sys.argv) < 2:
        sys.exit("Usage: python write_double_servo_test.py <goal_rad>")

    goal_rad = float(sys.argv[1])
    if not math.isfinite(goal_rad):
        sys.exit(f"Goal is not finite: {goal_rad!r}")

    goals_ticks = []
    for mid, sgn in zip(MOTOR_IDS, SIGNS):
        ticks = rad_to_ticks(goal_rad, sgn)
        if ticks < SAFE_TICK_LO or ticks > SAFE_TICK_HI:
            sys.exit(
                f"Goal {goal_rad:+.3f} rad = {ticks} ticks for motor {mid} "
                f"is outside safe range [{SAFE_TICK_LO}, {SAFE_TICK_HI}]"
            )
        goals_ticks.append(ticks)

    port = PortHandler(USB_PORT)
    if not port.openPort():
        sys.exit("Failed to open port")
    if not port.setBaudRate(BAUD_RATE):
        sys.exit("Failed to set baud rate")
    sts = sms_sts(port)

    try:
        print(f"Goal: {goal_rad:+.3f} rad -> {goals_ticks[0]} ticks (both motors)")

        present_ticks = []
        for mid in MOTOR_IDS:
            _, comm, err = sts.ping(mid)
            if comm != 0 or err != 0:
                sys.exit(f"Motor {mid} not responding (comm={comm}, err={err})")
            pos, _, _ = sts.ReadPos(mid)
            present_ticks.append(pos)
            rad = (pos - CENTER_TICK) * TICKS_TO_RAD
            print(f"  motor {mid}: pos={pos} ({rad:+.3f} rad)")

        # Safety handshake (simultaneous via RegWrite + RegAction)
        print("\nHandshake: writing current positions...")
        for mid, curr in zip(MOTOR_IDS, present_ticks):
            sts.RegWritePosEx(mid, int(curr), DEFAULT_SPEED, DEFAULT_ACCEL)
        sts.RegAction()
        time.sleep(0.2)

        # Write goal to both motors simultaneously
        print(f"Writing goal to both motors (simultaneous)...")
        for mid, goal in zip(MOTOR_IDS, goals_ticks):
            sts.RegWritePosEx(mid, int(goal), DEFAULT_SPEED, DEFAULT_ACCEL)
        sts.RegAction()

        # Monitor
        print("\nMonitoring... (ctrl-C to abort)")
        period = 1.0 / MONITOR_HZ
        frame = 0
        deadline = time.time() + MONITOR_TIMEOUT_S
        while time.time() < deadline:
            t0 = time.perf_counter()
            positions = []
            moving_any = False
            for mid in MOTOR_IDS:
                pos, _, _ = sts.ReadPos(mid)
                mov, _, _ = sts.ReadMoving(mid)
                positions.append(pos)
                if mov:
                    moving_any = True

            if frame % TEMP_POLL_EVERY_N == 0:
                for mid in MOTOR_IDS:
                    temp, comm, _ = sts.read1ByteTxRx(mid, ADDR_TEMP)
                    if comm != 0:
                        continue
                    if temp >= TEMP_ABORT_C:
                        print(f"\n!!! TEMP ABORT motor {mid} = {temp}C !!!")
                        emergency_stop(sts)
                        return
                    elif temp >= TEMP_WARN_C:
                        print(f"\n  warning: motor {mid} = {temp}C")

            diff = positions[0] - positions[1]
            r1 = SIGNS[0] * (positions[0] - CENTER_TICK) * TICKS_TO_RAD
            r2 = SIGNS[1] * (positions[1] - CENTER_TICK) * TICKS_TO_RAD
            status = "moving" if moving_any else "stopped"
            print(f"  m1={positions[0]:>5} ({r1:>+.3f})  m2={positions[1]:>5} ({r2:>+.3f})  "
                  f"diff={diff:>+4d}  [{status}]")

            if not moving_any and frame > 2:
                print("Both motors stopped.")
                break

            frame += 1
            dt = time.perf_counter() - t0
            if dt < period:
                time.sleep(period - dt)

        # Pin to present (simultaneous)
        for mid in MOTOR_IDS:
            pos, _, _ = sts.ReadPos(mid)
            sts.RegWritePosEx(mid, int(pos), DEFAULT_SPEED, DEFAULT_ACCEL)
        sts.RegAction()

        # Final
        print("\nFinal:")
        for mid in MOTOR_IDS:
            pos, _, _ = sts.ReadPos(mid)
            print(f"  motor {mid}: pos={pos} ({(pos - CENTER_TICK) * TICKS_TO_RAD:+.3f} rad)")

        try:
            if input("\nRelease torque? [y/N] ").strip().lower() == "y":
                emergency_stop(sts)
                print("Torque disabled.")
            else:
                print("Torque stays on.")
        except EOFError:
            pass

    except KeyboardInterrupt:
        print("\n[INTERRUPTED] disabling torque...")
        emergency_stop(sts)
    except Exception as e:
        print(f"\n[ERROR] {e!r} — disabling torque...")
        emergency_stop(sts)
        raise
    finally:
        port.closePort()


if __name__ == "__main__":
    main()