#!/usr/bin/env python3
"""
Hold the Inspire hand OPEN and probe ALL 12 joints at once — live position + force
+ deviation from a calibrated baseline. Press a finger and watch which joint(s)
react (and whether the force bleeds into neighbours). Logs everything to CSV.

  conda activate g1_env
  cd ~/Robotics_workspace/yslootahtech/G1/G1_Lootah/Manual_Recorder
  python3 hand_joint_probe.py enp3s0                 # hold open, live table + CSV
  python3 hand_joint_probe.py enp3s0 --thresh 40     # more sensitive press flag
  python3 hand_joint_probe.py enp3s0 --no-log        # view only

Needs inspire_g1 running on the robot (the Docker container) so it publishes
rt/inspire/state with force in tau_est. Ctrl+C to stop.
q: 1.00 open .. 0.00 closed. Idx 0-5 = right, 6-11 = left.
"""
import sys, time, argparse, csv
from datetime import datetime
from unitree_sdk2py.core.channel import ChannelPublisher, ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.unitree_go.msg.dds_ import MotorCmds_, MotorStates_
from unitree_sdk2py.idl.default import unitree_go_msg_dds__MotorCmd_

NAMES = ["R.pinky", "R.ring", "R.mid", "R.index", "R.thumbB", "R.thumbR",
         "L.pinky", "L.ring", "L.mid", "L.index", "L.thumbB", "L.thumbR"]
CLR = "\x1b[K"


def bar(v, w=16):
    v = 0.0 if v < 0 else (1.0 if v > 1 else v)
    n = int(round(v * w))
    return "█" * n + "·" * (w - n)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("iface", nargs="?", default="enp3s0")
    ap.add_argument("--thresh", type=float, default=30.0, help="press-flag threshold (g); auto-floored to per-joint noise")
    ap.add_argument("--hz", type=float, default=15.0)
    ap.add_argument("--log", default=None)
    ap.add_argument("--no-log", action="store_true")
    ap.add_argument("--no-hold", action="store_true", help="don't command open; just observe current pose")
    a = ap.parse_args()

    ChannelFactoryInitialize(0, a.iface)
    st = {"q": [0.0] * 12, "f": [0.0] * 12, "n": 0}

    def on_s(m):
        try:
            st["q"] = [float(m.states[i].q) for i in range(12)]
            st["f"] = [float(m.states[i].tau_est) for i in range(12)]
            st["n"] += 1
        except Exception:
            pass

    ChannelSubscriber("rt/inspire/state", MotorStates_).Init(on_s, 10)
    pub = ChannelPublisher("rt/inspire/cmd", MotorCmds_)
    pub.Init()
    msg = MotorCmds_()
    msg.cmds = [unitree_go_msg_dds__MotorCmd_() for _ in range(12)]

    def hold_open():
        if a.no_hold:
            return
        for i in range(12):
            msg.cmds[i].q = 1.0
            msg.cmds[i].kp = 0.0
        pub.Write(msg)

    # Hold open + calibrate the resting force baseline (and per-joint noise).
    print("Holding hand OPEN + calibrating force baseline — DON'T touch (2s)...")
    acc = [0.0] * 12
    fmin = [1e9] * 12
    fmax = [-1e9] * 12
    nc = 0
    t0 = time.time()
    while time.time() - t0 < 2.0:
        hold_open()
        if any(abs(v) > 1e-6 for v in st["f"]):
            for i in range(12):
                acc[i] += st["f"][i]
                fmin[i] = min(fmin[i], st["f"][i])
                fmax[i] = max(fmax[i], st["f"][i])
            nc += 1
        time.sleep(0.05)
    base = [acc[i] / nc if nc else 0.0 for i in range(12)]
    thr = [max(a.thresh, (fmax[i] - fmin[i]) * 1.5 + 10) if nc else a.thresh for i in range(12)]
    peak = [0.0] * 12
    if not nc:
        print("⚠️  no force data — is inspire_g1 up and the adapter connected?")

    writer = logf = logpath = None
    if not a.no_log:
        logpath = a.log or f"joint_probe_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        logf = open(logpath, "w", newline="")
        writer = csv.writer(logf)
        writer.writerow(["t_epoch", "iso"] +
                        [f"pos_{n}" for n in NAMES] +
                        [f"force_{n}" for n in NAMES] +
                        [f"dev_{n}" for n in NAMES] +      # Δbase
                        [f"peak_{n}" for n in NAMES] +     # peakΔ (running max)
                        [f"thr_{n}" for n in NAMES] +      # per-joint press threshold
                        [f"react_{n}" for n in NAMES])     # 1 = pressed this frame

    dt = 1.0 / a.hz
    print("\x1b[2J", end="")
    try:
        while True:
            now = time.time()
            hold_open()
            q, f = st["q"], st["f"]
            lines = ["\x1b[H"]
            lines.append(f"Inspire hand — ALL-joints probe {'(held OPEN)' if not a.no_hold else '(observe)'}"
                         f"   msgs={st['n']:>6d}   {datetime.now().strftime('%H:%M:%S')}   (Ctrl+C){CLR}")
            lines.append(f"press a finger → the reacting joint(s) light up.  pos: 1.00 open ░ 0.00 closed{CLR}")
            lines.append(CLR)
            lines.append(f"  {'joint':9s} {'position':20s} {'force':>7s} {'Δbase':>7s} {'peakΔ':>6s}  react{CLR}")
            for i, nm in enumerate(NAMES):
                if i == 6:
                    lines.append(CLR)
                dev = f[i] - base[i]
                peak[i] = max(peak[i], abs(dev))
                hit = abs(dev) > thr[i]
                pos = f"[{bar(q[i])}] {q[i]:4.2f}"
                bold = "\x1b[1m" if hit else ""
                mark = "◄ PRESS" if hit else ""
                lines.append(f"  {bold}{nm:9s} {pos} {f[i]:6.0f}g {dev:+6.0f} {peak[i]:6.0f}  {mark}\x1b[0m{CLR}")
            lines.append(CLR)
            lr = any(abs(x) > 1e-3 for x in f[:6])
            ll = any(abs(x) > 1e-3 for x in f[6:])
            lines.append(f"  RIGHT {'● live' if lr else '○ dropped'}     LEFT {'● live' if ll else '○ dropped'}"
                         + (f"     log: {logpath}" if logpath else "") + CLR)
            print("\n".join(lines), end="", flush=True)
            if writer:
                devs = [f[i] - base[i] for i in range(12)]
                writer.writerow([f"{now:.3f}", datetime.now().isoformat()] +
                                [f"{x:.4f}" for x in q] +
                                [f"{x:.1f}" for x in f] +
                                [f"{d:.1f}" for d in devs] +
                                [f"{peak[i]:.1f}" for i in range(12)] +
                                [f"{thr[i]:.1f}" for i in range(12)] +
                                [1 if abs(devs[i]) > thr[i] else 0 for i in range(12)])
            time.sleep(dt)
    except KeyboardInterrupt:
        pass
    finally:
        print("\n\n  peak Δ per joint (g)   [per-joint threshold]:")
        for i, nm in enumerate(NAMES):
            mark = "✓ pressed" if peak[i] > thr[i] else ""
            print(f"    {nm:9s} {peak[i]:6.0f}  (thr {thr[i]:5.0f})  {mark}")
        if logf:
            logf.close()
        print("\nstopped." + (f"  log: {logpath}" if logpath else ""))


if __name__ == "__main__":
    main()
