"""Write joint goals to Feetech servos with safety rails.

Usage:
  python write_motors_from_goals.py <motor_id> [motor_id ...]
    [--goals r1,r2,...] [--speed N] [--accel N] [--yes]

Without --goals, uses the built-in DEFAULT_GOALS_RAD truncated to the
number of motor IDs passed. Default speed and accel are intentionally
very slow (10/10) so you can watch and kill power if something looks
wrong.

Safety layers (in order):
  1. Pings every motor before doing anything.
  2. Clamps each goal to SAFE_TICK_RANGE and aborts if any goal is
     non-finite or outside the range.
  3. Writes *current* positions as the initial goal (handshake), so
     nothing snaps if the goal register had a stale value.
  4. Uses WritePosEx with slow speed + accel (servo-side trapezoid).
  5. Loops on ReadMoving until all motors stop (with a hard timeout).
  6. Polls temperature every ~0.5 s; warns at TEMP_WARN_C, aborts at
     TEMP_ABORT_C (disables torque and exits).
  7. SIGINT / any exception -> emergency_stop (torque off all motors)
     then close the port.
"""
import argparse
import math
import sys
import time

import mujoco
import mujoco.viewer

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

DEFAULT_XML = (
    "/Users/cameronsmith/Projects/robotics_testing/"
    "menagerie_testing/mujoco_menagerie/2ourso100/example_twolink.xml"
)

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 window per joint: tighter than the servo's hardware limits.
# ±1251 ticks = ±1.91986 rad = the model's joint range.
SAFE_TICK_LO = CENTER_TICK - 1251  # 797
SAFE_TICK_HI = CENTER_TICK + 1251  # 3299

# Normal defaults — matches the "fast" mode from so100_controller.
DEFAULT_SPEED = 360  # ticks / s
DEFAULT_ACCEL = 120  # ticks / s^2

ADDR_TORQUE_ENABLE = 40
ADDR_TEMP = 63
TEMP_WARN_C = 55
TEMP_ABORT_C = 65
TEMP_POLL_EVERY_N_FRAMES = 10  # ~0.5 s at 20 Hz

MONITOR_HZ = 20
MONITOR_TIMEOUT_S = 120.0

# Taken from the MuJoCo viewer pose the user grabbed.
# Order matches the MJCF joint order: shoulder_pan..shoulder_pan7.
DEFAULT_GOALS_RAD = [-0.11, -0.38, 0.03, -0.12, 0.09, 0.0, 0.0]


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


def ticks_to_rad_signed(ticks: int, sign: int) -> float:
    return sign * (ticks - CENTER_TICK) * TICKS_TO_RAD


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("motor_ids", nargs="+", type=int)
    p.add_argument("--goals", default=None,
                   help="Comma-separated goal radians, one per motor.")
    p.add_argument("--signs", default=None,
                   help="Comma-separated +1/-1 per motor (default all +1).")
    p.add_argument("--speed", type=int, default=DEFAULT_SPEED)
    p.add_argument("--accel", type=int, default=DEFAULT_ACCEL)
    p.add_argument("--xml", default=DEFAULT_XML,
                   help="MJCF to load for live visualization.")
    return p.parse_args()


def collect_hinge_joints(model):
    """Return (name, qpos_adr) for each hinge joint, in XML order."""
    joints = []
    for i in range(model.njnt):
        if model.jnt_type[i] == mujoco.mjtJoint.mjJNT_HINGE:
            name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)
            joints.append((name, int(model.jnt_qposadr[i])))
    return joints


def emergency_stop(sts, motor_ids):
    """Disable torque on every motor. Uses writeTxRx (not write1ByteTxRx)
    to avoid the scservo_patch portHandler-double-injection bug."""
    for mid in motor_ids:
        try:
            sts.writeTxRx(mid, ADDR_TORQUE_ENABLE, 1, [0])
        except Exception as e:
            print(f"  [estop] failed to torque-off motor {mid}: {e}")


def read_temp(sts, mid):
    val, comm, _ = sts.read1ByteTxRx(mid, ADDR_TEMP)
    if comm != 0:
        return None
    return val


def check_temperatures(sts, motor_ids):
    """Return (abort, any_warn). Prints warnings/alerts as a side effect."""
    abort = False
    any_warn = False
    for mid in motor_ids:
        temp = read_temp(sts, mid)
        if temp is None:
            continue
        if temp >= TEMP_ABORT_C:
            print(f"!!! TEMP ABORT: motor {mid} = {temp}C >= {TEMP_ABORT_C}C !!!")
            abort = True
        elif temp >= TEMP_WARN_C:
            print(f"  warning: motor {mid} = {temp}C >= {TEMP_WARN_C}C")
            any_warn = True
    return abort, any_warn


def main():
    args = parse_args()
    n = len(args.motor_ids)

    if args.goals:
        cleaned = args.goals.replace("[", " ").replace("]", " ").replace(",", " ")
        goals_rad = [float(x) for x in cleaned.split()]
    else:
        if n > len(DEFAULT_GOALS_RAD):
            sys.exit(
                f"DEFAULT_GOALS_RAD has {len(DEFAULT_GOALS_RAD)} values, "
                f"but you passed {n} motor IDs"
            )
        goals_rad = DEFAULT_GOALS_RAD[:n]

    if len(goals_rad) != n:
        sys.exit(f"{len(goals_rad)} goal values but {n} motor IDs")

    if args.signs:
        signs = [int(s) for s in args.signs.split(",")]
        if len(signs) != n:
            sys.exit("--signs length must match number of motor IDs")
    else:
        signs = [1] * n

    # Validate goals: finite + within safe tick window.
    goals_ticks = []
    for mid, rad, sgn in zip(args.motor_ids, goals_rad, signs):
        if not math.isfinite(rad):
            sys.exit(f"Goal for motor {mid} is not finite: {rad!r}")
        ticks = rad_to_ticks(rad, sgn)
        if ticks < SAFE_TICK_LO or ticks > SAFE_TICK_HI:
            sys.exit(
                f"Goal for motor {mid} = {ticks} ticks ({rad:+.3f} rad) "
                f"is outside safe range [{SAFE_TICK_LO}, {SAFE_TICK_HI}]"
            )
        goals_ticks.append(ticks)

    # Load MuJoCo model for live viz. Map motor IDs to the first N hinge joints.
    model = mujoco.MjModel.from_xml_path(args.xml)
    data = mujoco.MjData(model)
    all_joints = collect_hinge_joints(model)
    if n > len(all_joints):
        sys.exit(
            f"Got {n} motor IDs but model only has {len(all_joints)} hinge joints"
        )
    joints = all_joints[:n]
    unmapped = all_joints[n:]
    if unmapped:
        print(f"Ignoring unmapped joints: {[j for j, _ in unmapped]}")

    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)

    viewer = mujoco.viewer.launch_passive(model, data)

    try:
        print("Pinging motors and reading state...")
        present_ticks = []
        for mid in args.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)
            temp = read_temp(sts, mid)
            present_ticks.append(pos)
            rad = (pos - CENTER_TICK) * TICKS_TO_RAD
            temp_str = f"{temp}C" if temp is not None else "?"
            if temp is not None and temp >= TEMP_ABORT_C:
                sys.exit(f"Motor {mid} temperature {temp}C >= {TEMP_ABORT_C}C. Aborting.")
            if temp is not None and temp >= TEMP_WARN_C:
                print(f"  warning: motor {mid} temperature {temp}C is elevated")
            print(f"  motor {mid}: pos={pos} ({rad:+.3f} rad)  temp={temp_str}")

        # Plan table
        print("\nPlanned moves:")
        header = f"{'motor':>6} {'curr':>6} {'goal':>6} {'delta':>7} {'goal rad':>10}"
        print(header)
        for mid, curr, goal, rad in zip(args.motor_ids, present_ticks, goals_ticks, goals_rad):
            print(f"{mid:>6} {curr:>6} {goal:>6} {goal - curr:>+7} {rad:>+10.3f}")
        print(f"profile: speed={args.speed} ticks/s, accel={args.accel} ticks/s^2")
        print(f"safe tick range: [{SAFE_TICK_LO}, {SAFE_TICK_HI}]")

        # Seed viewer with current positions.
        for (_, adr), pos in zip(joints, present_ticks):
            data.qpos[adr] = (pos - CENTER_TICK) * TICKS_TO_RAD
        mujoco.mj_forward(model, data)
        viewer.sync()

        # --- Safety handshake: write current positions as the initial goal.
        print("\n[1/2] Writing current positions as initial goal (no motion expected)...")
        for mid, curr in zip(args.motor_ids, present_ticks):
            sts.RegWritePosEx(mid, int(curr), args.speed, args.accel)
        sts.RegAction()
        time.sleep(0.2)

        # --- Real goal (all motors start simultaneously)
        print("[2/2] Writing target goals...")
        for mid, goal in zip(args.motor_ids, goals_ticks):
            sts.RegWritePosEx(mid, int(goal), args.speed, args.accel)
        sts.RegAction()

        # --- Monitor until stopped (or timeout / temp abort / ctrl-C)
        print("\nMonitoring motion... (ctrl-C to abort and disable torque)")
        period = 1.0 / MONITOR_HZ
        frame = 0
        deadline = time.time() + MONITOR_TIMEOUT_S
        temp_abort = False
        while time.time() < deadline and viewer.is_running():
            t0 = time.perf_counter()
            states = []
            moving_any = False
            for (name, adr), mid, sgn in zip(joints, args.motor_ids, signs):
                pos, _, _ = sts.ReadPos(mid)
                mov, _, _ = sts.ReadMoving(mid)
                states.append((pos, mov))
                if mov:
                    moving_any = True
                data.qpos[adr] = sgn * (pos - CENTER_TICK) * TICKS_TO_RAD
            mujoco.mj_forward(model, data)
            viewer.sync()

            if frame % TEMP_POLL_EVERY_N_FRAMES == 0:
                temp_abort, _ = check_temperatures(sts, args.motor_ids)
                if temp_abort:
                    break

            row = "  ".join(
                f"m{mid}:{pos:>4}{'M' if mov else 'S'}"
                for mid, (pos, mov) in zip(args.motor_ids, states)
            )
            print(f"  t={frame * period:4.1f}s  {row}")

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

            frame += 1
            dt = time.perf_counter() - t0
            if dt < period:
                time.sleep(period - dt)
        else:
            print("Timeout waiting for motors to stop.")

        # Pin each motor's goal to its actual present position so the PID
        # stops integrating a tiny residual error and any dead-zone /
        # backlash hunting settles. Cheap, reversible.
        print("\nPinning goals to present positions to suppress residual hunting...")
        settled_ticks = []
        for mid in args.motor_ids:
            pos, _, _ = sts.ReadPos(mid)
            settled_ticks.append(pos)
            sts.RegWritePosEx(mid, int(pos), args.speed, args.accel)
        sts.RegAction()

        # Final state
        print("\nFinal positions:")
        for mid, pos in zip(args.motor_ids, settled_ticks):
            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, args.motor_ids)
                print("Torque disabled on all motors.")
            else:
                print("Leaving torque enabled so the arm holds position.")
        except EOFError:
            print("Leaving torque enabled so the arm holds position.")

        # Keep viewer live-updating with real servo state until closed.
        print("\nLive view (close the MuJoCo window to exit)...")
        idle_period = 1.0 / MONITOR_HZ
        idle_frame = 0
        while viewer.is_running():
            t0 = time.perf_counter()
            for (_, adr), mid, sgn in zip(joints, args.motor_ids, signs):
                pos, comm, _ = sts.ReadPos(mid)
                if comm != 0:
                    continue
                data.qpos[adr] = sgn * (pos - CENTER_TICK) * TICKS_TO_RAD
            mujoco.mj_forward(model, data)
            viewer.sync()
            if idle_frame % TEMP_POLL_EVERY_N_FRAMES == 0:
                check_temperatures(sts, args.motor_ids)
            idle_frame += 1
            dt = time.perf_counter() - t0
            if dt < idle_period:
                time.sleep(idle_period - dt)

    except KeyboardInterrupt:
        print("\n[INTERRUPTED] disabling torque on all motors...")
        emergency_stop(sts, args.motor_ids)
    except Exception as e:
        print(f"\n[ERROR] {e!r} — disabling torque on all motors...")
        emergency_stop(sts, args.motor_ids)
        raise
    finally:
        try:
            viewer.close()
        except Exception:
            pass
        port.closePort()


if __name__ == "__main__":
    main()