"""Live-visualize Feetech servo positions in the MuJoCo viewer.

Reads each servo's position over serial and writes it into the matching
joint's qpos, then syncs the passive viewer at ~30Hz.

Usage:
  python read_and_vis_motors.py <motor_id> [motor_id ...]

Motor IDs are mapped in order to the model's hinge joints as they appear
in the MJCF (shoulder_pan, shoulder_pan2, ...). Pass one ID per joint.

Options:
  --xml PATH      Path to the MJCF file to load.
  --signs S       Comma-separated +1/-1 per motor, in the same order as
                  motor IDs. Default: all +1.
  --rate HZ       Poll rate (default 30).
"""
import argparse
import math
import sys
import time
import numpy as np

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

ADDR_TEMP = 63
TEMP_WARN_C = 55  # print warning at/above this
TEMP_ALERT_C = 65  # print loud alert at/above this (read-only: we don't auto-stop)
TEMP_POLL_EVERY_N_FRAMES = 30  # at 30Hz this is ~1 Hz


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("motor_ids", nargs="+", type=int)
    p.add_argument("--xml", default=DEFAULT_XML)
    p.add_argument("--signs", default=None,
                   help="Comma-separated +1/-1 per motor (default all +1)")
    p.add_argument("--rate", type=float, default=30.0)
    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 ticks_to_rad(pos_ticks, sign):
    return sign * (pos_ticks - CENTER_TICK) * TICKS_TO_RAD


def main():
    args = parse_args()

    model = mujoco.MjModel.from_xml_path(args.xml)
    data = mujoco.MjData(model)

    all_joints = collect_hinge_joints(model)
    if len(args.motor_ids) > len(all_joints):
        sys.exit(
            f"Got {len(args.motor_ids)} motor IDs but model only has "
            f"{len(all_joints)} hinge joints: {[n for n, _ in all_joints]}"
        )
    joints = all_joints[: len(args.motor_ids)]
    unmapped = all_joints[len(args.motor_ids):]
    if unmapped:
        print(f"Ignoring unmapped joints: {[n for n, _ in unmapped]}")

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

    print("Joint -> motor mapping:")
    for (name, _), mid, sgn in zip(joints, args.motor_ids, signs):
        print(f"  {name:16s} <- motor {mid} (sign {sgn:+d})")

    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)

    for mid in args.motor_ids:
        _, comm, err = sts.ping(mid)
        if comm != 0 or err != 0:
            port.closePort()
            sys.exit(f"Motor {mid} not responding (comm={comm}, err={err})")

    period = 1.0 / args.rate
    frame = 0
    try:
        with mujoco.viewer.launch_passive(model, data) as viewer:
            while viewer.is_running():
                t0 = time.perf_counter()
                for (name, adr), mid, sgn in zip(joints, args.motor_ids, signs):
                    pos, comm, _ = sts.ReadPos(mid)
                    if comm != 0:
                        continue  # transient; keep last value
                    data.qpos[adr] = ticks_to_rad(pos, sgn)
                mujoco.mj_forward(model, data)
                viewer.sync()
                if frame % TEMP_POLL_EVERY_N_FRAMES == 0:
                    for (name, _), mid in zip(joints, args.motor_ids):
                        temp, comm, _ = sts.read1ByteTxRx(mid, ADDR_TEMP)
                        if comm != 0:
                            continue
                        if temp >= TEMP_ALERT_C:
                            print(f"!!! TEMP ALERT: motor {mid} ({name}) = {temp}C "
                                  f">= {TEMP_ALERT_C}C — consider stopping !!!")
                        elif temp >= TEMP_WARN_C:
                            print(f"  warning: motor {mid} ({name}) = {temp}C "
                                  f">= {TEMP_WARN_C}C")
                dt = time.perf_counter() - t0
                if dt < period:
                    time.sleep(period - dt)
                print(",".join(str(x) for x in np.array(data.qpos).round(2)))
                frame += 1
    finally:
        port.closePort()


if __name__ == "__main__":
    main()