#!/usr/bin/env python3
"""
Hand-ONLY record + replay for the Inspire hand. Ways to record:

  --mode press   (default) : TOUCH a finger and it gets recorded — even a SLIGHT touch.
                             Detection is on Δbase (force change from rest, threshold
                             ~12g vs a ~5g noise floor); any detected touch closes the
                             finger at least --min-close (35%) and HOLDS it --hold (0.8s)
                             so a quick tap is clearly visible on replay. Touch harder =
                             closes more (fully closed at --span, 45g). The fingers stay
                             OPEN while recording and only close on REPLAY.
  --mode levels            : each press steps OPEN → SLIGHT → HALF → FULL → OPEN.
  --mode release           : fingers RELEASED (loose) — pose by hand, records positions.
                             NOTE: the encoder does not see hand-posing on this hand, so
                             this usually records nothing — prefer press mode.
  --mode capture           : pose by hand and tap SPACE to snapshot each pose.

  conda activate g1_env
  cd ~/Robotics_workspace/yslootahtech/G1/G1_Lootah/Manual_Recorder
  python3 hand_record.py enp3s0 --output grip --seconds 20    # touch fingers -> recorded
  python3 hand_record.py enp3s0 --replay grip                 # replay a saved file
  # more sensitive still:  --thresh 8 --span 30 --min-close 0.5 --hold 1.2

Keys during recording: o = re-open all · Ctrl+C = end early.
Needs inspire_g1 running (Docker container). Files go in DataHand/.
q: 1.00 = open, 0.00 = fully closed.  Idx 0-5 = right hand, 6-11 = left.
"""
import sys, time, json, argparse
from pathlib import Path
import termios, tty, select
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_

HAND_MOTORS = 12
REC_HZ = 30.0
GRIP_FORCE = 500.0
DATA_DIR = Path("DataHand")
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"
# Discrete close levels: each press of a finger steps to the next one.
# (name, commanded position) — 1.00 = open, 0.00 = fully closed.
LEVELS = [("OPEN", 1.00), ("SLIGHT", 0.75), ("HALF", 0.50), ("FULL", 0.05)]


class KeyPoller:
    def __enter__(self):
        self.fd = sys.stdin.fileno()
        self.old = termios.tcgetattr(self.fd)
        tty.setcbreak(self.fd)
        return self
    def __exit__(self, *a):
        termios.tcsetattr(self.fd, termios.TCSADRAIN, self.old)
    def poll(self):
        if select.select([sys.stdin], [], [], 0)[0]:
            return sys.stdin.read(1)
        return None


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)


class HandRecorder:
    def __init__(self, iface):
        ChannelFactoryInitialize(0, iface)
        self.st = {"q": [1.0] * 12, "f": [0.0] * 12, "n": 0}
        ChannelSubscriber("rt/inspire/state", MotorStates_).Init(self._on_state, 10)
        self.pub = ChannelPublisher("rt/inspire/cmd", MotorCmds_)
        self.pub.Init()
        self.msg = MotorCmds_()
        self.msg.cmds = [unitree_go_msg_dds__MotorCmd_() for _ in range(12)]
        time.sleep(0.5)

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

    def send(self, q, force=GRIP_FORCE):
        for i in range(12):
            self.msg.cmds[i].q = float(q[i])
            self.msg.cmds[i].kp = float(force)
        self.pub.Write(self.msg)

    def release(self):
        self.send([-1.0] * 12, force=0.0)     # motor OFF -> fingers loose

    def open_hand(self, n=25):
        for _ in range(n):
            self.send([1.0] * 12)
            time.sleep(0.02)

    # ---------- mode: release (pose by hand) ----------
    def record_release(self, seconds):
        print("🖐️  Opening hand, then releasing so you can pose it by hand...")
        self.open_hand()
        for _ in range(10):
            self.release(); time.sleep(0.02)
        print("\x1b[2J", end="")
        frames, dt, t0 = [], 1.0 / REC_HZ, time.time()
        try:
            with KeyPoller() as kp:
                while time.time() - t0 < seconds:
                    loop = time.time()
                    k = kp.poll()
                    if k and k.lower() == "o":
                        self.open_hand(10)
                    self.release()
                    frames.append({"t": round(time.time() - t0, 4), "h": list(self.st["q"])})
                    self._draw("REC (pose by hand)", self.st["q"])
                    time.sleep(max(0, dt - (time.time() - loop)))
        except KeyboardInterrupt:
            pass
        print(f"\n\n✅ recorded {len(frames)} frames ({frames[-1]['t'] if frames else 0:.1f}s).")
        return frames

    # ---------- mode: capture (snapshot poses by hand) ----------
    def record_capture(self, hold_sec):
        print("🖐️  Opening + releasing. Pose the fingers by hand, tap SPACE to CAPTURE each")
        print("    pose you like, Ctrl+C when done. Replay walks through the captured poses.")
        self.open_hand()
        for _ in range(10):
            self.release(); time.sleep(0.02)
        keyframes, flash = [], 0
        print("\x1b[2J", end="")
        try:
            with KeyPoller() as kp:
                while True:
                    loop = time.time()
                    k = kp.poll()
                    if k == " ":
                        keyframes.append(list(self.st["q"]))
                        flash = 8
                    elif k and k.lower() == "o":
                        self.open_hand(10)
                    self.release()
                    tag = f"CAPTURE — {len(keyframes)} pose(s) saved" + ("   ✓ CAPTURED!" if flash else "")
                    self._draw(tag + "   (SPACE=capture · o=open · Ctrl+C=done)", self.st["q"])
                    flash = max(0, flash - 1)
                    time.sleep(max(0, 1.0 / REC_HZ - (time.time() - loop)))
        except KeyboardInterrupt:
            pass
        print(f"\n\n✅ captured {len(keyframes)} pose(s).")
        # Build a time-series: hold each captured pose for hold_sec (motor ramps between).
        frames, t, dt, n = [], 0.0, 1.0 / REC_HZ, int(hold_sec * REC_HZ)
        for kf in keyframes:
            for _ in range(n):
                frames.append({"t": round(t, 4), "h": list(kf)})
                t += dt
        return frames

    # ---------- mode: press (proportional to Δbase, recorded + replayed) ----------
    def record_press(self, seconds, thresh, span, signed, min_close=0.35, hold=0.8, auto_scale=True):
        """Continuously monitor each finger's Δbase (force change from rest) and map it
        to how far that finger closes. Records BOTH the raw Δbase track and the mapped
        positions; replay drives the fingers through the mapped positions."""
        print("🤛 Opening hand + calibrating Δbase — DON'T touch (2s)...")
        self.open_hand()
        acc = [0.0] * 12; fmin = [1e9] * 12; fmax = [-1e9] * 12; nc = 0
        t0 = time.time()
        while time.time() - t0 < 2.0:
            self.send([1.0] * 12)
            if any(abs(v) > 1e-6 for v in self.st["f"]):
                for i in range(12):
                    acc[i] += self.st["f"][i]
                    fmin[i] = min(fmin[i], self.st["f"][i]); fmax[i] = max(fmax[i], self.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(thresh, (fmax[i] - fmin[i]) * 1.5 + 5) if nc else thresh for i in range(12)]
        if nc:
            print("   thresholds R:", [round(thr[i]) for i in range(6)],
                  " L:", [round(thr[i]) for i in range(6, 12)])
        else:
            print("   ⚠️  no force data — inspire_g1 up? adapter connected?")

        peak = [0.0] * 12
        latch = [0.0] * 12          # held closeness per finger
        latch_until = [0.0] * 12    # hold the latch until this time
        touches = [0] * 12          # how many touches were captured per finger
        touching = [False] * 12
        print("\x1b[2J", end="")
        frames, dt, t0 = [], 1.0 / REC_HZ, time.time()
        try:
            with KeyPoller() as kp:
                while time.time() - t0 < seconds:
                    loop = time.time()
                    kp.poll()
                    now = time.time()
                    devs, vals = [0.0] * 12, [1.0] * 12
                    for i in range(12):
                        d = self.st["f"][i] - base[i]
                        devs[i] = d
                        peak[i] = max(peak[i], abs(d))
                        # sign: negative = pushed toward closing. --signed uses that
                        # direction; default uses magnitude (either sign = a press).
                        amt = (-d if signed else abs(d))
                        closeness = (amt - thr[i]) / max(1.0, span - thr[i])
                        closeness = min(1.0, max(0.0, closeness))
                        if amt > thr[i]:
                            # ANY detected touch closes at least min_close, so even the
                            # slightest tap is clearly visible on replay.
                            closeness = max(closeness, min_close)
                            if not touching[i]:
                                touches[i] += 1
                                touching[i] = True
                            latch[i] = max(latch[i], closeness)
                            latch_until[i] = now + hold      # keep holding while touched
                        elif amt < thr[i] * 0.5:
                            touching[i] = False
                        # latch: hold the closure briefly after the touch ends, then release
                        if now < latch_until[i]:
                            closeness = max(closeness, latch[i])
                        else:
                            latch[i] = 0.0
                        vals[i] = 1.0 - closeness          # 1.0 open .. 0.0 closed
                    self.send([1.0] * 12)                  # finger stays OPEN while recording
                    frames.append({"t": round(time.time() - t0, 4),
                                   "h": list(vals), "d": [round(x, 1) for x in devs]})
                    self._draw_press(f"REC (Δbase → close) {seconds-(time.time()-t0):4.1f}s left",
                                     devs, vals, thr, peak)
                    time.sleep(max(0, dt - (time.time() - loop)))
        except KeyboardInterrupt:
            pass
        print(f"\n\n✅ recorded {len(frames)} frames ({frames[-1]['t'] if frames else 0:.1f}s).")
        # Re-derive positions from the Δbase track so the DEGREE of each touch survives.
        frames, pk = self._remap(frames, thr, min_close, hold, auto_scale, span)
        tot = sum(touches)
        print(f"   {tot} touch(es) captured"
              + (" — each finger scaled to its own hardest press:" if auto_scale else ":")
              if tot else
              "   ⚠️  no touches captured — lower --thresh (try 8) and touch the fingertip pad.")
        for i in range(12):
            if touches[i] or peak[i] > thr[i] * 0.5:
                closes = [1.0 - f["h"][i] for f in frames]
                mx = max(closes) if closes else 0.0
                if touches[i]:
                    print(f"     {NAMES[i]:9s} peakΔ {pk[i]:5.0f}g  {touches[i]} touch(es)"
                          f"  → closes up to {mx*100:3.0f}% on replay")
                else:
                    print(f"     {NAMES[i]:9s} peakΔ {pk[i]:5.0f}g  (below thr {thr[i]:.0f}, not captured)")
        return frames

    def _remap(self, frames, thr, min_close, hold, auto_scale, span):
        """Re-derive the replay positions from the recorded Δbase track so the DEGREE of
        each touch is preserved. With auto_scale, every finger is normalised to its OWN
        hardest press (that press = fully closed, lighter ones proportionally less), which
        matters because a thumb reaches ~190g while an index may only reach ~50g."""
        if not frames:
            return frames, [0.0] * 12
        pk = [0.0] * 12
        for f in frames:
            for i in range(12):
                pk[i] = max(pk[i], abs(f["d"][i]))
        for f in frames:
            for i in range(12):
                amt = abs(f["d"][i])
                if amt <= thr[i]:
                    c = 0.0
                else:
                    top = pk[i] if (auto_scale and pk[i] > thr[i] + 5) else span
                    top = max(top, thr[i] + 5)
                    frac = min(1.0, (amt - thr[i]) / (top - thr[i]))
                    # threshold touch -> min_close, hardest touch -> fully closed
                    c = min_close + (1.0 - min_close) * frac
                f["h"][i] = 1.0 - c
        # latch: hold each touch's closure briefly so a quick tap stays visible
        w = int(hold * REC_HZ)
        for i in range(12):
            cl = [1.0 - f["h"][i] for f in frames]
            run, cnt = 0.0, 0
            for k in range(len(frames)):
                if cl[k] > 0:
                    run, cnt = max(run, cl[k]), w
                elif cnt > 0:
                    cnt -= 1
                else:
                    run = 0.0
                frames[k]["h"][i] = 1.0 - max(cl[k], run if cnt > 0 else 0.0)
        return frames, pk

    def _draw_press(self, tag, devs, vals, thr, peak):
        out = ["\x1b[H", f"HAND RECORD — {tag}   msgs={self.st['n']:>6d}{CLR}",
               "  press a finger: the harder you press, the more it will close on replay" + CLR,
               f"  {'joint':9s} {'Δbase':>7s} {'peakΔ':>6s} {'thr':>4s}  close-on-replay{CLR}", CLR]
        for i, nm in enumerate(NAMES):
            if i == 6:
                out.append(CLR)
            hit = "◄" if abs(devs[i]) > thr[i] else " "
            out.append(f"  {nm:9s} {devs[i]:+7.0f} {peak[i]:6.0f} {thr[i]:4.0f} {hit} "
                       f"[{bar(1.0 - vals[i])}] {(1.0-vals[i])*100:3.0f}%{CLR}")
        print("\n".join(out), end="", flush=True)

    # ---------- mode: force (press joints, based on Δbase) ----------
    def record_force(self, seconds, thresh, close_rate):
        print("🤛 Opening hand + calibrating Δbase — DON'T touch (2s)...")
        self.open_hand()
        acc = [0.0] * 12; fmin = [1e9] * 12; fmax = [-1e9] * 12; nc = 0
        t0 = time.time()
        while time.time() - t0 < 2.0:
            self.send([1.0] * 12)
            if any(abs(v) > 1e-6 for v in self.st["f"]):
                for i in range(12):
                    acc[i] += self.st["f"][i]
                    fmin[i] = min(fmin[i], self.st["f"][i]); fmax[i] = max(fmax[i], self.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(thresh, (fmax[i] - fmin[i]) * 1.5 + 10) if nc else thresh for i in range(12)]
        lvl = [0] * 12                       # level index per finger
        pushed = [False] * 12
        flash = [0] * 12
        print("\x1b[2J", end="")
        frames, dt, t0 = [], 1.0 / REC_HZ, time.time()
        try:
            with KeyPoller() as kp:
                while time.time() - t0 < seconds:
                    loop = time.time()
                    k = kp.poll()
                    if k:
                        kl = k.lower()
                        if kl == "o": lvl = [0] * 12
                        elif kl == "c": lvl = [3] * 12
                        elif kl == "[":
                            for i in range(6): lvl[i] = 0
                        elif kl == "]":
                            for i in range(6): lvl[i] = 3
                        elif kl == ";":
                            for i in range(6, 12): lvl[i] = 0
                        elif kl == "'":
                            for i in range(6, 12): lvl[i] = 3
                    for i in range(12):
                        dev = self.st["f"][i] - base[i]
                        if abs(dev) > thr[i]:                     # press detected via Δbase
                            if not pushed[i]:                     # rising edge -> step one level
                                lvl[i] = (lvl[i] + 1) % len(LEVELS)
                                pushed[i] = True
                                flash[i] = 12
                        elif abs(dev) < thr[i] * 0.5:             # released (hysteresis)
                            pushed[i] = False
                    self.send([1.0] * 12)                         # finger stays OPEN while recording
                    vals = [LEVELS[lvl[i]][1] for i in range(12)]
                    frames.append({"t": round(time.time() - t0, 4), "h": list(vals)})
                    self._draw_levels(f"REC (press a finger = next level) {seconds-(time.time()-t0):4.1f}s",
                                      lvl, flash)
                    for i in range(12):
                        flash[i] = max(0, flash[i] - 1)
                    time.sleep(max(0, dt - (time.time() - loop)))
        except KeyboardInterrupt:
            pass
        print(f"\n\n✅ recorded {len(frames)} frames ({frames[-1]['t'] if frames else 0:.1f}s). "
              f"Fingers stayed open; they close to these levels on replay:")
        for i in range(12):
            if lvl[i]:
                print(f"    {NAMES[i]:9s} {LEVELS[lvl[i]][0]}")
        return frames

    def _draw_levels(self, tag, lvl, flash):
        out = ["\x1b[H", f"HAND RECORD — {tag}   msgs={self.st['n']:>6d}{CLR}",
               "  press a finger → OPEN → SLIGHT → HALF → FULL → OPEN   (no need to close it!)" + CLR,
               "  keys: o/c all open/full · [ ] right · ; ' left · Ctrl+C=end" + CLR, CLR]
        for i, nm in enumerate(NAMES):
            if i == 6:
                out.append(CLR)
            name, val = LEVELS[lvl[i]]
            mark = " ← set" if flash[i] else ""
            out.append(f"  {nm:9s} [{bar(val)}] {name:6s}{mark}{CLR}")
        print("\n".join(out), end="", flush=True)

    def _draw(self, tag, vals, label="position"):
        out = ["\x1b[H", f"HAND RECORD — {tag}   msgs={self.st['n']:>6d}   (o=open all · Ctrl+C=end){CLR}",
               f"  {label}: 1.00 open ░ 0.00 closed" + CLR, CLR]
        for i, nm in enumerate(NAMES):
            if i == 6:
                out.append(CLR)
            out.append(f"  {nm:9s} [{bar(vals[i])}] {vals[i]:4.2f}{CLR}")
        print("\n".join(out), end="", flush=True)

    def replay(self, frames, speed=1.0):
        if not frames:
            print("   (nothing to replay)")
            return
        print("▶️  replaying — the motor drives the fingers through the recording...")
        first = frames[0]["h"]
        cur = list(self.st["q"])
        for k in range(20):
            a = k / 20.0
            self.send([(1 - a) * cur[i] + a * first[i] for i in range(12)])
            time.sleep(1.0 / REC_HZ)
        print("\x1b[2J", end="")
        t0, i = time.time(), 0
        try:
            while True:
                tgt = (time.time() - t0) * speed
                while i < len(frames) - 1 and frames[i]["t"] < tgt:
                    i += 1
                self.send(frames[i]["h"])
                self._draw("REPLAY", frames[i]["h"])
                if i >= len(frames) - 1 and frames[-1]["t"] < tgt:
                    break
                time.sleep(1.0 / REC_HZ)
            print("\n\n✅ replay done.")
        except KeyboardInterrupt:
            print("\n\n⛔ replay stopped.")


def save(frames, name):
    DATA_DIR.mkdir(exist_ok=True)
    p = DATA_DIR / (name if name.endswith(".jsonl") else name + ".jsonl")
    with open(p, "w") as f:
        f.write(json.dumps({"meta": {"hz": REC_HZ, "hand_motors": HAND_MOTORS}}) + "\n")
        for fr in frames:
            f.write(json.dumps(fr) + "\n")
    print(f"💾 saved {len(frames)} frames to {p}")


def load(name):
    p = DATA_DIR / (name if name.endswith(".jsonl") else name + ".jsonl")
    frames = []
    with open(p) as f:
        f.readline()
        for line in f:
            d = json.loads(line)
            if "h" in d:
                frames.append(d)
    print(f"loaded {len(frames)} frames from {p}")
    return frames


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("iface", nargs="?", default="enp3s0")
    ap.add_argument("--output", default="hand_rec")
    ap.add_argument("--seconds", type=float, default=20.0)
    ap.add_argument("--mode", choices=["press", "levels", "release", "capture", "force"], default="press",
                    help="press=Δbase proportional — harder press closes the finger more, recorded "
                         "continuously and replayed (default); levels=press steps OPEN/SLIGHT/HALF/FULL; "
                         "release=pose by hand; capture=snapshot poses (SPACE); force=alias of levels")
    ap.add_argument("--span", type=float, default=45.0,
                    help="press mode: Δbase (g) mapping to a FULLY closed finger; lower = lighter touch closes more (default 45)")
    ap.add_argument("--min-close", type=float, default=0.35,
                    help="press mode: ANY detected touch closes at least this much, 0-1 (default 0.35)")
    ap.add_argument("--hold", type=float, default=0.8,
                    help="press mode: seconds to hold a touch's closure so a quick tap is clearly replayed (default 0.8)")
    ap.add_argument("--no-auto-scale", action="store_true",
                    help="press mode: use the fixed --span instead of scaling each finger to its own hardest press")
    ap.add_argument("--signed", action="store_true",
                    help="press mode: use force DIRECTION (negative Δbase = closing) instead of magnitude")
    ap.add_argument("--pose-hold", type=float, default=1.5, help="capture mode: seconds to hold each pose on replay")
    ap.add_argument("--replay", default=None, help="replay a saved hand file (name in DataHand/) and exit")
    ap.add_argument("--speed", type=float, default=1.0)
    ap.add_argument("--thresh", type=float, default=12.0,
                    help="Δbase press threshold in grams. Noise floor is ~5g, so 12 catches a SLIGHT touch (default 12)")
    ap.add_argument("--close-rate", type=float, default=1.2, help="force mode: degree ramp speed (per s)")
    a = ap.parse_args()

    rec = HandRecorder(a.iface)
    if a.replay:
        rec.replay(load(a.replay), a.speed)
        return

    if a.mode == "press":
        frames = rec.record_press(a.seconds, a.thresh, a.span, a.signed, a.min_close, a.hold,
                                  not a.no_auto_scale)
    elif a.mode in ("levels", "force"):
        frames = rec.record_force(a.seconds, a.thresh, a.close_rate)
    elif a.mode == "capture":
        frames = rec.record_capture(a.pose_hold)
    else:
        frames = rec.record_release(a.seconds)
    if input("Replay the recording? [y/n]: ").strip().lower() in ("y", "yes", ""):
        rec.replay(frames, a.speed)
    if input("Save recording? [y/n]: ").strip().lower() in ("y", "yes", ""):
        save(frames, a.output)
    print("done.")


if __name__ == "__main__":
    main()