#!/usr/bin/env python3
"""
g1_damp_live_joint_movement.py

- Tries to put Unitree G1 into Damp mode (requires Loco RPC service / lease).
- ALWAYS runs live LowState joint monitoring (does NOT require LocoClient).

If LocoClient.Init() fails, it falls back to MONITOR-ONLY mode and still
shows which joints are moving when you move them by hand.

Usage:
  python3 g1_damp_live_joint_movement.py --iface enp3s0
  python3 g1_damp_live_joint_movement.py --iface enp3s0 --no-damp
"""

from __future__ import annotations

import argparse
import math
import sys
import time
from typing import Dict, List, Optional, Tuple

from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelSubscriber
from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_  # type: ignore


# ---------------- Joint map (your table) ----------------
JOINTS: List[Tuple[int, str, Tuple[float, float]]] = [
    (0,  "L_LEG_HIP_PITCH",   (-2.5307, 2.8798)),
    (1,  "L_LEG_HIP_ROLL",    (-0.5236, 2.9671)),
    (2,  "L_LEG_HIP_YAW",     (-2.7576, 2.7576)),
    (3,  "L_LEG_KNEE",        (-0.087267, 2.8798)),
    (4,  "L_LEG_ANKLE_PITCH", (-0.87267, 0.5236)),
    (5,  "L_LEG_ANKLE_ROLL",  (-0.2618, 0.2618)),
    (6,  "R_LEG_HIP_PITCH",   (-2.5307, 2.8798)),
    (7,  "R_LEG_HIP_ROLL",    (-2.9671, 0.5236)),
    (8,  "R_LEG_HIP_YAW",     (-2.7576, 2.7576)),
    (9,  "R_LEG_KNEE",        (-0.087267, 2.8798)),
    (10, "R_LEG_ANKLE_PITCH", (-0.87267, 0.5236)),
    (11, "R_LEG_ANKLE_ROLL",  (-0.2618, 0.2618)),
    (12, "WAIST_YAW",         (-2.618, 2.618)),
    (13, "WAIST_ROLL",        (-0.52, 0.52)),
    (14, "WAIST_PITCH",       (-0.52, 0.52)),
    (15, "L_SHOULDER_PITCH",  (-3.0892, 2.6704)),
    (16, "L_SHOULDER_ROLL",   (-1.5882, 2.2515)),
    (17, "L_SHOULDER_YAW",    (-2.618, 2.618)),
    (18, "L_ELBOW",           (-1.0472, 2.0944)),
    (19, "L_WRIST_ROLL",      (-1.972222054, 1.972222054)),
    (20, "L_WRIST_PITCH",     (-1.614429558, 1.614429558)),
    (21, "L_WRIST_YAW",       (-1.614429558, 1.614429558)),
    (22, "R_SHOULDER_PITCH",  (-3.0892, 2.6704)),
    (23, "R_SHOULDER_ROLL",   (-2.2515, 1.5882)),
    (24, "R_SHOULDER_YAW",    (-2.618, 2.618)),
    (25, "R_ELBOW",           (-1.0472, 2.0944)),
    (26, "R_WRIST_ROLL",      (-1.972222054, 1.972222054)),
    (27, "R_WRIST_PITCH",     (-1.614429558, 1.614429558)),
    (28, "R_WRIST_YAW",       (-1.614429558, 1.614429558)),
]
JOINT_BY_IDX: Dict[int, Tuple[str, Tuple[float, float]]] = {i: (n, lim) for i, n, lim in JOINTS}


def rad2deg(x: float) -> float:
    return float(x) * 180.0 / math.pi


def parse_only(s: str) -> set[int]:
    out: set[int] = set()
    for part in s.split(","):
        part = part.strip()
        if part:
            out.add(int(part))
    return out


def clear_screen() -> None:
    sys.stdout.write("\x1b[2J\x1b[H")


def safe_call(name: str, fn, *args, **kwargs):
    try:
        return fn(*args, **kwargs)
    except Exception as e:
        print(f"[ERR] {name} failed: {e}", file=sys.stderr)
        return None


class LiveMovementUI:
    def __init__(self, only: Optional[set[int]], hz: float, move_threshold_deg: float):
        self.only = only
        self.period = 1.0 / max(hz, 1e-6)
        self.move_thr = float(move_threshold_deg)

        self.last_print = 0.0
        self.msg_count = 0

        self.prev_q_rad: Dict[int, float] = {}
        self.last_q_rad: Dict[int, float] = {}
        self.last_movers: List[Tuple[int, float]] = []

    def cb(self, msg: LowState_):
        self.msg_count += 1
        motor_state = getattr(msg, "motor_state", None)
        if motor_state is None:
            return

        for idx in JOINT_BY_IDX.keys():
            try:
                self.last_q_rad[idx] = float(motor_state[idx].q)
            except Exception:
                continue

        now = time.time()
        if now - self.last_print < self.period:
            return
        self.last_print = now

        movers: List[Tuple[int, float]] = []
        for idx, q in self.last_q_rad.items():
            if self.only is not None and idx not in self.only:
                continue
            q_prev = self.prev_q_rad.get(idx, q)
            ddeg = abs(rad2deg(q - q_prev))
            if ddeg >= self.move_thr:
                movers.append((idx, ddeg))

        movers.sort(key=lambda x: x[1], reverse=True)
        self.last_movers = movers[:10]
        self.prev_q_rad = dict(self.last_q_rad)

        self.render()

    def render(self):
        clear_screen()
        sys.stdout.write(
            f"G1 Live Joint Monitor | msgs={self.msg_count} | "
            f"print_hz≈{1.0/self.period:.1f} | move_thr={self.move_thr:.2f}°\n"
        )
        sys.stdout.write("Move joints by hand; moving joints are marked with '>>'\n")
        if self.only:
            sys.stdout.write(f"Filter: only {sorted(self.only)}\n")
        sys.stdout.write("-" * 100 + "\n")

        if self.last_movers:
            sys.stdout.write("Most moving now:\n")
            for idx, ddeg in self.last_movers:
                name, _ = JOINT_BY_IDX.get(idx, ("UNKNOWN", (0.0, 0.0)))
                qdeg = rad2deg(self.last_q_rad.get(idx, 0.0))
                sys.stdout.write(f"  >> [{idx:02d}] {name:<18}  Δ={ddeg:6.2f}°   q={qdeg:+8.2f}°\n")
            sys.stdout.write("-" * 100 + "\n")
        else:
            sys.stdout.write("Most moving now: (none above threshold)\n")
            sys.stdout.write("-" * 100 + "\n")

        sys.stdout.write(f"{'IDX':>3}  {'JOINT':<20}  {'q(deg)':>10}  {'MOVING':>8}\n")
        sys.stdout.write("-" * 100 + "\n")

        for idx in sorted(JOINT_BY_IDX.keys()):
            if self.only is not None and idx not in self.only:
                continue

            name, _lim = JOINT_BY_IDX[idx]
            qrad = self.last_q_rad.get(idx, 0.0)
            qdeg = rad2deg(qrad)

            q_prev = self.prev_q_rad.get(idx, qrad)
            ddeg = abs(rad2deg(qrad - q_prev))
            moving = ">>" if ddeg >= self.move_thr else ""

            sys.stdout.write(f"{idx:>3}  {name:<20}  {qdeg:+10.2f}  {moving:>8}\n")

        sys.stdout.flush()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--iface", required=True, help="NIC connected to robot (e.g. enp3s0)")
    ap.add_argument("--topic", default="", help="Override topic (default tries rt/lf/lowstate then rt/lowstate)")
    ap.add_argument("--hz", type=float, default=10.0, help="Print rate (Hz). Default 10")
    ap.add_argument("--move-threshold-deg", type=float, default=0.25, help="Mark joint moving if Δdeg >= this")
    ap.add_argument("--only", default="", help='Comma indices, e.g. "15,16,17,18,19,20,21"')
    ap.add_argument("--timeout", type=float, default=10.0)
    ap.add_argument("--no-damp", action="store_true", help="Skip Damp (monitor only)")
    args = ap.parse_args()

    only = parse_only(args.only) if args.only else None

    # DDS init
    ChannelFactoryInitialize(0, args.iface)

    # Try Damp (optional)
    bot_ok = False
    if not args.no_damp:
        bot = LocoClient()
        safe_call("SetTimeout", bot.SetTimeout, float(args.timeout))
        if safe_call("Init", bot.Init) is not None:
            bot_ok = True
            safe_call("Damp", bot.Damp)
            print("[OK] Sent Damp(). You can move joints by hand.")
        else:
            print(
                "[WARN] LocoClient.Init() failed (lease/service not available). "
                "Continuing in MONITOR-ONLY mode.\n"
                "Tip: this is the same root issue as your lease error 3102."
            )
    else:
        print("[INFO] --no-damp set: running MONITOR-ONLY (no control calls).")

    ui = LiveMovementUI(only=only, hz=args.hz, move_threshold_deg=args.move_threshold_deg)

    topics = [args.topic] if args.topic else ["rt/lf/lowstate", "rt/lowstate"]

    sub = None
    last_err = None
    for t in topics:
        try:
            sub = ChannelSubscriber(t, LowState_)
            sub.Init(ui.cb, 200)
            print(f"[OK] Subscribed: {t}")
            break
        except Exception as e:
            last_err = e
            print(f"[WARN] Subscribe failed for {t}: {e}", file=sys.stderr)

    if sub is None:
        raise SystemExit(f"[FATAL] Could not subscribe to LowState. Last error: {last_err}")

    if not bot_ok and not args.no_damp:
        print("\n[NOTE] Robot may NOT be in Damp because control service init failed.")
        print("      If you need compliance, set Damp from robot UI/remote, then monitor here.\n")

    print("Running... Ctrl+C to stop.")
    try:
        while True:
            time.sleep(1.0)
    except KeyboardInterrupt:
        print("\nStopping.")
        # If we had bot control, leave it in Damp
        try:
            if not args.no_damp and bot_ok:
                safe_call("Damp", bot.Damp)
        except Exception:
            pass


if __name__ == "__main__":
    main()
