"""Back-drive calibrated Feetech servos and mirror them live in the MuJoCo viewer.

Drives only the servos that are BOTH calibrated (in calib/<config>.json) AND currently
responding on the bus — so a partially-connected/partially-calibrated arm just works,
with a note on what's live vs missing. Torque is disabled so you can move the joints by
hand and watch the sim follow. Read-only: it never commands the servos.

Usage:
  python live_mirror.py --config arm [--port /dev/tty...] [--rate 40] [--model PATH]
"""
import argparse
import json
import os
import sys
import time

import mujoco
import mujoco.viewer

from feetech_bus import Bus, ticks_to_rad
from configs import CONFIGS, calib_path

TEMP_WARN_C = 60


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="gripper")
    ap.add_argument("--port", default=None)
    ap.add_argument("--rate", type=float, default=40.0)
    ap.add_argument("--model", default=None, help="Override the config's MuJoCo model path")
    a = ap.parse_args()

    cfg = CONFIGS[a.config]
    model_path = a.model or cfg["model"]
    path = calib_path(a.config)
    try:
        calib = json.load(open(path))
    except FileNotFoundError:
        sys.exit(f"No calibration at {path}. Run calibrate_servo.py first.")
    if not calib:
        sys.exit(f"{path} is empty. Run calibrate_servo.py first.")

    model = mujoco.MjModel.from_xml_path(model_path)
    data = mujoco.MjData(model)

    # Resolve each calibrated servo to a model joint (skip calib joints not in the model).
    entries = []          # (servo_id, qpos_adr, sign, joint)
    not_in_model = []
    for sid, c in sorted(calib.items(), key=lambda kv: int(kv[0])):
        jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, c["joint"])
        if jid < 0:
            not_in_model.append((sid, c["joint"]))
            continue
        entries.append((int(sid), int(model.jnt_qposadr[jid]), int(c["sign"]), c["joint"]))
    if not entries:
        sys.exit(f"None of the calibrated joints exist in {os.path.basename(model_path)}.")

    # Partition by what's actually on the bus right now.
    bus = Bus(a.port) if a.port else Bus()
    live, missing = [], []
    for e in entries:
        (live if bus.ping(e[0]) else missing).append(e)

    all_hinges = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)
                  for i in range(model.njnt) if int(model.jnt_type[i]) == mujoco.mjtJoint.mjJNT_HINGE]
    covered = {e[3] for e in entries}
    uncalibrated = [j for j in all_hinges if j not in covered]

    print(f"Model: {os.path.basename(model_path)}")
    print(f"LIVE + calibrated ({len(live)}) — driven by hand:")
    for sid, _, sgn, jn in live:
        print(f"    servo {sid:>2} -> {jn}  (sign {sgn:+d})")
    if missing:
        print(f"MISSING ({len(missing)}) — calibrated but not responding, held at rest:")
        for sid, _, _, jn in missing:
            print(f"    servo {sid:>2} -> {jn}")
    if uncalibrated:
        print(f"NOT CALIBRATED — held at rest: {uncalibrated}")
    if not_in_model:
        print(f"IN CALIB BUT NOT IN MODEL — ignored: {not_in_model}")
    if not live:
        bus.close()
        sys.exit("No calibrated servos are responding — nothing to mirror.")

    for sid, _, _, _ in live:
        bus.torque(sid, False)  # back-drivable

    period = 1.0 / a.rate
    frame = 0
    try:
        with mujoco.viewer.launch_passive(model, data) as viewer:
            print("\nMove a live joint by hand — the sim follows. Close the viewer or Ctrl-C to stop.")
            while viewer.is_running():
                t0 = time.perf_counter()
                for sid, adr, sgn, _ in live:
                    pos = bus.read_pos(sid)
                    if pos is None:
                        continue  # transient; keep last value
                    data.qpos[adr] = ticks_to_rad(pos, sgn)
                mujoco.mj_forward(model, data)
                viewer.sync()
                if frame % int(max(1, a.rate)) == 0:  # ~1 Hz temp check
                    for sid, _, _, jn in live:
                        tmp = bus.read_temp(sid)
                        if tmp is not None and tmp >= TEMP_WARN_C:
                            print(f"  warning: servo {sid} ({jn}) = {tmp}C")
                frame += 1
                dt = time.perf_counter() - t0
                if dt < period:
                    time.sleep(period - dt)
    finally:
        bus.close()


if __name__ == "__main__":
    main()