#!/usr/bin/env python3
"""
FORCE-FOLLOW — shape the Inspire hand by pushing its fingers, like posing a limp hand.

The RH56 cannot be back-driven: its worm drive self-locks, so pushing a finger moves it
0.000 rad no matter how hard you press (measured: 233 g of push, 0.000-0.005 of rotation).
Kinesthetic teaching by physically bending the fingers is therefore impossible.

This makes it FEEL back-drivable instead, with admittance control — the standard technique
for exactly this class of device (non-backdrivable, high-friction, force-sensed):

    push the pad  ->  force sensor reads it  ->  motor drives the finger that way
    stop pushing  ->  it holds where it is

So you shape the hand with your hands, and the pose sticks. Then record it.

    conda activate g1_env
    cd ~/Robotics_workspace/yslootahtech/G1/G1_Lootah/Manual_Recorder
    python3 hand_follow.py enp3s0                  # push fingers to close them
    python3 hand_follow.py enp3s0 --invert         # if a push OPENS instead of closes
    python3 hand_follow.py enp3s0 --gain 0.0009    # faster / slower response

Ctrl-C prints the shape you ended on, ready to paste into the dashboard or a recording.

SAFETY: a finger you push closes WHILE you keep pushing. If it closes onto your fingertip
the force stays high and it keeps going, so it is speed-limited (full travel ~3s) and the
hand's own grip-force limit still applies. Push with a knuckle or the back of a pen if you
would rather not have it close on your skin.
"""
import sys, time, argparse
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"]
THUMB_ROT = [5, 11]

ap = argparse.ArgumentParser()
ap.add_argument("iface", nargs="?", default="enp3s0")
ap.add_argument("--gain", type=float, default=0.0006,
                help="travel per gram per second. higher = more eager (default 0.0006)")
ap.add_argument("--deadband", type=float, default=30.0,
                help="grams of push before a finger reacts, ignores drift/noise")
ap.add_argument("--maxrate", type=float, default=0.35,
                help="max travel per second (0.35 = full range in ~3s)")
ap.add_argument("--force", type=float, default=400.0, help="grip force limit while shaping")
ap.add_argument("--invert", action="store_true", help="flip which way a push moves the finger")
ap.add_argument("--only", default="", help="comma list to enable, e.g. R.index,R.mid (default all)")
a = ap.parse_args()

enabled = set(range(12))
if a.only.strip():
    enabled = set()
    for tok in a.only.split(","):
        tok = tok.strip().lower()
        hit = [i for i, n in enumerate(NAMES) if n.lower() == tok]
        if not hit:
            sys.exit(f"--only: unknown finger {tok!r}. Names: {', '.join(NAMES)}")
        enabled.add(hit[0])

ChannelFactoryInitialize(0, a.iface)
st = {"q": [1.0] * 12, "f": [0.0] * 12, "n": 0}
def on_state(m):
    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
ChannelSubscriber("rt/inspire/state", MotorStates_).Init(on_state, 10)
pub = ChannelPublisher("rt/inspire/cmd", MotorCmds_); pub.Init()
msg = MotorCmds_(); msg.cmds = [unitree_go_msg_dds__MotorCmd_() for _ in range(12)]

def send(q):
    for i in range(12):
        msg.cmds[i].q = float(max(0.0, min(1.0, q[i])))
        msg.cmds[i].kp = a.force
    pub.Write(msg)

time.sleep(1.0)
if st["n"] == 0:
    sys.exit("no hand state — is inspire_g1 running?")

# Start from wherever the hand already is, so it does not jump when you begin.
target = list(st["q"])
send(target)

print("calibrating resting force — DON'T touch (2s)...")
acc = [0.0] * 12; n = 0
t0 = time.time()
while time.time() - t0 < 2.0:
    send(target)
    for i in range(12): acc[i] += st["f"][i]
    n += 1
    time.sleep(0.03)
base = [acc[i] / n for i in range(12)]

print("\n" + "=" * 66)
print("PUSH A FINGER — it follows while you push, and holds when you stop")
print(f"  gain {a.gain}  deadband {a.deadband:.0f}g  max {a.maxrate}/s"
      + ("  INVERTED" if a.invert else ""))
if a.only.strip():
    print("  enabled: " + ", ".join(NAMES[i] for i in sorted(enabled)))
print("  Ctrl-C to stop and print the shape")
print("=" * 66 + "\n")

sign = 1.0 if a.invert else -1.0     # push on the pad (+force) should CLOSE (q toward 0)
last = time.time()
try:
    while True:
        now = time.time(); dt = now - last; last = now
        dt = min(dt, 0.1)
        moving = []
        for i in range(12):
            if i not in enabled:
                continue
            dev = st["f"][i] - base[i]
            # Measured: idle force drift over 90s untouched is 0g on all 12 fingers, so the
            # baseline is stable and needs no re-centring. The deadband only has to reject
            # incidental brushes, which the integrator would otherwise accumulate.
            if abs(dev) <= a.deadband:
                continue
            # push harder -> move faster, but never faster than maxrate
            step = sign * a.gain * (dev - (a.deadband if dev > 0 else -a.deadband)) * dt
            step = max(-a.maxrate * dt, min(a.maxrate * dt, step))
            target[i] = max(0.0, min(1.0, target[i] + step))
            moving.append(f"{NAMES[i].split('.')[1][:3]}{dev:+.0f}g→{target[i]:.2f}")
        send(target)
        # ALWAYS show the live force on every finger, not just ones over the deadband. Three
        # separate background measurements read 0g simply because they were not running while
        # a push happened; a live readout removes that whole class of confusion — you can see
        # immediately whether a push registers, how many grams it is, and which sign.
        top = sorted(((abs(st["f"][i] - base[i]), i) for i in range(12)), reverse=True)[:6]
        live = "  ".join(
            f"{NAMES[i].split('.')[1][:3]}{st['f'][i]-base[i]:+.0f}"
            + ("*" if abs(st["f"][i]-base[i]) > a.deadband else "")
            for v, i in top if v > 1)
        act = " ".join(moving)
        print("\r  " + (f"MOVING {act}" if act else f"force: {live or 'all quiet'}   "
              f"(need >{a.deadband:.0f}g)")[:110] + "\x1b[K", end="", flush=True)
        time.sleep(0.03)
except KeyboardInterrupt:
    print("\n\nshape you ended on:")
    print("  " + "  ".join(f"{NAMES[i].split('.')[1][:3]}={st['q'][i]:.2f}" for i in range(6)))
    print("  " + "  ".join(f"{NAMES[i].split('.')[1][:3]}={st['q'][i]:.2f}" for i in range(6, 12)))
    print("\n  as 12 values (paste into the dashboard / a pose file):")
    print("  " + " ".join(f"{st['q'][i]:.3f}" for i in range(12)))
    print("\n  holding this shape. Ctrl-C again to exit.")
    try:
        while True:
            send(target); time.sleep(0.05)
    except KeyboardInterrupt:
        print("done.")