#!/usr/bin/env python3
"""
G1 TEACH MODE V4 (STABLE & PROVEN)
----------------------------------
1. NETWORK: 60Hz (Safe for WiFi/Ethernet).
2. WAIST: LOCKED (Kp=60). Keeps robot rigid and stable.
3. ARMS: Safe Start (3s) -> Limp Teach -> Return Home.

Usage:
  python3 g1_teach_v4_stable.py enp3s0 --output high_five --seconds 15
  python3 g1_teach_v4_stable.py enp3s0 --output photo_G2 --seconds 25 --home arm_home.jsonl
  python3 g1_teach_v4_stable.py enp3s0 --output adnoc --seconds 25 --home arm_home.jsonl
"""

import time
import sys
import json
import argparse
import numpy as np
import os
import termios
import tty
import select
from pathlib import Path
from unitree_sdk2py.core.channel import ChannelPublisher, ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowState_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_
from unitree_sdk2py.utils.crc import CRC
# Hand (Inspire) command + state IDL — same lib, unitree_go
from unitree_sdk2py.idl.unitree_go.msg.dds_ import MotorCmds_, MotorStates_
from unitree_sdk2py.idl.default import unitree_go_msg_dds__MotorCmd_

G1_NUM_MOTOR = 29
ENABLE_ARM_SDK_INDEX = 29
WAIST_INDICES = [12, 13, 14]
ARM_INDICES = list(range(15, 29))
CONTROL_INDICES = WAIST_INDICES + ARM_INDICES
DATA_DIR = Path("DataG1")
HAND_TOPIC = "rt/inspire/cmd"   # Inspire hand command stream (state read-back is unreliable)
HAND_MOTORS = 12                # 0-5 right, 6-11 left

RECORD_HZ = 60.0
SAFE_START_TIME = 3.0

# Tuned Gains (The "Amazing" Settings)
WAIST_KP = 60.0
WAIST_KD = 4.0

HOLD_ARM_KP = 60.0
HOLD_ARM_KD = 4.0
TEACH_ARM_KP = 0.0
TEACH_ARM_KD = 2.0
# Preview/replay-before-save drives the arms firmer so they reach the SAME height as
# the real replay (g1_replay uses KP_LOW=80); Kp=60 sags lower under gravity.
PREVIEW_ARM_KP = 80.0
PREVIEW_ARM_KD = 3.0

YES_CHOICES = {"yes", "y", "1", "t", "true"}
NO_CHOICES = {"no", "n", "0", "f", "false"}

def ask_yes_no(prompt: str):
    while True:
        try:
            answer = input(prompt).strip().lower()
        except KeyboardInterrupt:
            print("\n⛔ Ctrl+C pressed.")
            return None
        if answer in YES_CHOICES:
            return True
        if answer in NO_CHOICES:
            return False
        print("Please answer yes or no.")

def ensure_unique_path(user_path: str) -> str:
    p = Path(user_path)
    if not p.suffix: p = p.with_suffix(".jsonl")
    if len(p.parts) > 1:
        target_dir = p.parent
        target_dir.mkdir(parents=True, exist_ok=True)
        candidate = target_dir / p.name
    else:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        candidate = DATA_DIR / p.name

    if not candidate.exists(): return str(candidate)
    i = 1
    while True:
        cand2 = candidate.with_name(f"{candidate.stem}({i}){candidate.suffix}")
        if not cand2.exists(): return str(cand2)
        i += 1

# --- SAME CONCEPT AS REPLAY: resolve_input_path + load_home_pose ---
def resolve_input_path(in_path: str) -> str:
    p = Path(in_path)
    if len(p.parts) == 1:
        return str(DATA_DIR / p.name)
    return str(p)

def load_home_pose(home_path: str):
    """Reads the last frame of arm_home.jsonl to get the target pose."""
    path = resolve_input_path(home_path)
    try:
        last_valid_q = None
        with open(path, 'r') as f:
            for line in f:
                d = json.loads(line)
                if 'q' in d and len(d['q']) == G1_NUM_MOTOR:
                    last_valid_q = d['q']
        if last_valid_q:
            print(f"✅ Loaded Home Pose from {path}")
            return last_valid_q
        else:
            print(f"⚠️  Warning: {path} found but contained no valid 'q' data.")
    except FileNotFoundError:
        print(f"⚠️  Warning: Home file {path} not found.")

    # Fallback: Zero Arms (same as replay fallback style)
    print("⚠️  Using Default Home (Arms at 0.0)")
    default_q = [0.0] * G1_NUM_MOTOR
    return default_q

def should_save_recording():
    return ask_yes_no("Save recording to file? [y/n]: ")

def should_repeat_recording():
    return ask_yes_no("Repeat recording? [y/n]: ")

# Non-blocking keyboard reader (set hand pose live during teach)
class KeyPoller:
    def __enter__(self):
        self.fd = sys.stdin.fileno()
        self.old_settings = termios.tcgetattr(self.fd)
        tty.setcbreak(self.fd)
        return self
    def __exit__(self, exc_type, exc, tb):
        termios.tcsetattr(self.fd, termios.TCSADRAIN, self.old_settings)
    def poll(self):
        if select.select([sys.stdin], [], [], 0)[0]:
            return sys.stdin.read(1)
        return None

# Hand presets for keyboard control during teach (q in [0,1])
HAND_PRESETS = {
    "o": [1.0] * HAND_MOTORS,                                   # open
    "c": [0.0] * HAND_MOTORS,                                   # close (fist)
    "b": [0.5, 0.5, 0.5, 0.5, 0.4, 0.1] * 2,                    # box
    "t": [0.0, 0.0, 0.0, 0.0, 1.0, 1.0] * 2,                    # thumbs-up
}

class StableTeacher:
    def __init__(self):
        self.low_state = None
        self.low_cmd = unitree_hg_msg_dds__LowCmd_()
        self.crc = CRC()
        self.arm_pub = ChannelPublisher("rt/arm_sdk", LowCmd_)
        self.arm_pub.Init()
        self.state_sub = ChannelSubscriber("rt/lowstate", LowState_)
        self.state_sub.Init(self.LowStateHandler, 10)
        self.first_state = False
        # Hand: record source = "state" (monitor actual finger degrees, default),
        # "keys" (keyboard presets), or "cmd" (dashboard command stream).
        self.record_hand = True
        self.hand_source = "keys"   # RH56 fingers are non-backdrivable; keys is the reliable way
        self.follow_force = 50.0                # low grip force for by-hand 'follow' (grams)
        self.default_grip = 500.0              # normal grip force to restore afterwards
        self.hand_q = [1.0] * HAND_MOTORS      # command / keyboard value
        self.hand_state = [0.0] * HAND_MOTORS  # latest actual finger positions (rt/inspire/state)
        # push-to-teach: press a finger -> force spikes (non-backdrivable, so position
        # can't move) -> we detect it and DRIVE that finger (toggle open/close).
        self.hand_force = [0.0] * HAND_MOTORS   # latest finger force (grams, from tau_est)
        self.hand_target = [1.0] * HAND_MOTORS  # commanded target in push mode
        self.force_base = [0.0] * HAND_MOTORS   # resting-force baseline (calibrated)
        self.pushed_prev = [False] * HAND_MOTORS
        self.force_peak = [0.0] * HAND_MOTORS   # peak |deviation| seen per finger (tuning aid)
        self.push_thresh = 50.0                 # grams deviation = a press (sensitive; light touch)
        self.push_thresh_eff = [50.0] * HAND_MOTORS  # per-finger, floored to resting noise at calib
        # push behaviour: "hold" = press-and-hold ramps the finger gradually (level of
        # close = how long/hard you hold), release holds; "toggle" = press flips full open/close.
        self.push_mode = "hold"
        self.close_rate = 1.2                   # finger units per second while held (1.0=open..0.0=closed)
        self.push_dir = [0.0] * HAND_MOTORS     # ramp direction chosen at each press start
        self.hand_sub = ChannelSubscriber(HAND_TOPIC, MotorCmds_)
        self.hand_sub.Init(self.HandHandler, 10)
        self.hand_state_sub = ChannelSubscriber("rt/inspire/state", MotorStates_)
        self.hand_state_sub.Init(self.HandStateHandler, 10)
        self.hand_pub = ChannelPublisher(HAND_TOPIC, MotorCmds_)
        self.hand_pub.Init()
        self.hand_msg = MotorCmds_()
        self.hand_msg.cmds = [unitree_go_msg_dds__MotorCmd_() for _ in range(HAND_MOTORS)]

    def HandStateHandler(self, msg: MotorStates_):
        try:
            self.hand_state = [float(msg.states[i].q) for i in range(HAND_MOTORS)]
            self.hand_force = [float(msg.states[i].tau_est) for i in range(HAND_MOTORS)]
        except Exception:
            pass

    def SendHand(self, h, force=0.0):
        # force>0 is relayed via the cmd kp field -> inspire_g1 sets the grip
        # (torque) force for that hand. A LOW force lets the fingers be pushed by
        # hand (motor can't resist) — used by 'follow' mode for by-hand recording.
        for i in range(HAND_MOTORS):
            self.hand_msg.cmds[i].q = float(h[i])
            self.hand_msg.cmds[i].kp = float(force)
        self.hand_pub.Write(self.hand_msg)

    def hand_frame(self):
        # what to record this frame, per source
        if self.hand_source == "push":
            return list(self.hand_target)  # commanded targets we drive the fingers to
        if self.hand_source in ("state", "follow"):
            return list(self.hand_state)   # actual sensed finger positions
        return list(self.hand_q)           # commanded (keys / dashboard)

    def LowStateHandler(self, msg: LowState_):
        self.low_state = msg
        self.first_state = True

    def HandHandler(self, msg: MotorCmds_):
        try:
            self.hand_q = [float(msg.cmds[i].q) for i in range(HAND_MOTORS)]
        except Exception:
            pass

    def SendCommand(self, q_targets, arm_kp, arm_kd, waist_kp, waist_kd):
        self.low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 1.0

        # Waist
        for i in WAIST_INDICES:
            self.low_cmd.motor_cmd[i].mode = 1
            self.low_cmd.motor_cmd[i].q = q_targets[i]
            self.low_cmd.motor_cmd[i].dq = 0
            self.low_cmd.motor_cmd[i].kp = waist_kp
            self.low_cmd.motor_cmd[i].kd = waist_kd
            self.low_cmd.motor_cmd[i].tau = 0

        # Arms
        for i in ARM_INDICES:
            self.low_cmd.motor_cmd[i].mode = 1
            self.low_cmd.motor_cmd[i].q = q_targets[i]
            self.low_cmd.motor_cmd[i].dq = 0
            self.low_cmd.motor_cmd[i].kp = arm_kp
            self.low_cmd.motor_cmd[i].kd = arm_kd
            self.low_cmd.motor_cmd[i].tau = 0

        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        self.arm_pub.Write(self.low_cmd)

    def DisableSDK(self):
        print("🔌 Disabling SDK... Handing control to internal brain.")
        self.low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 0.0
        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        interrupted = False
        for _ in range(5):
            self.arm_pub.Write(self.low_cmd)
            try:
                time.sleep(0.02)
            except KeyboardInterrupt:
                if not interrupted:
                    print("\n⚠️  Ctrl+C received. Finishing SDK disable...")
                    interrupted = True

    def ReturnToHome(self, initial_q, home_q, home_filename: str):
        print(f"\n🏡 FINISHED. RETURNING ARMS TO HOME ({home_filename})...")
        end_q = [self.low_state.motor_state[i].q for i in range(G1_NUM_MOTOR)]

        home_steps = 180  # ~3 seconds at 60Hz (same style as replay)
        interrupted = False
        for k in range(home_steps):
            alpha = (k / home_steps)

            interp_q = list(end_q)

            # Keep waist rigid at the original standing pose
            for j in WAIST_INDICES:
                interp_q[j] = initial_q[j]

            # Arms: interpolate End -> Home
            for j in ARM_INDICES:
                interp_q[j] = (1-alpha)*end_q[j] + alpha*home_q[j]

            try:
                self.SendCommand(interp_q, HOLD_ARM_KP, HOLD_ARM_KD, WAIST_KP, WAIST_KD)
                time.sleep(1.0 / RECORD_HZ)
            except KeyboardInterrupt:
                if not interrupted:
                    print("\n⚠️  Ctrl+C received. Continuing until home is reached...")
                    interrupted = True

        print("✅ Home Reached.")

    def Preview(self, samples, initial_q, home_q, home_filename: str):
        """Replay the just-recorded motion (arms + hand) so it can be previewed
        BEFORE deciding to save. SDK must still be enabled (call before DisableSDK).
        Moves to the first frame, plays in real time, then returns arms to home."""
        arm_samples = [s for s in samples
                       if "q" in s and len(s["q"]) == G1_NUM_MOTOR]
        if not arm_samples:
            print("   (nothing to preview)")
            return
        print("\n▶️  PREVIEW — replaying this recording (arms + hand will MOVE)...")

        # 1) Ease from the current (home) pose to the recording's first frame
        start_q = arm_samples[0]["q"]
        cur_q = [self.low_state.motor_state[i].q for i in range(G1_NUM_MOTOR)]
        move_steps = 120  # ~2s at 60Hz
        first_h = arm_samples[0].get("h")
        try:
            for k in range(move_steps):
                alpha = k / move_steps
                interp = list(cur_q)
                for j in WAIST_INDICES:
                    interp[j] = initial_q[j]           # waist locked to standing pose
                for j in ARM_INDICES:
                    interp[j] = (1 - alpha) * cur_q[j] + alpha * start_q[j]
                self.SendCommand(interp, PREVIEW_ARM_KP, PREVIEW_ARM_KD, WAIST_KP, WAIST_KD)
                if self.record_hand and first_h is not None:
                    self.SendHand(first_h)
                time.sleep(1.0 / RECORD_HZ)

            # 2) Play the recorded frames in real time
            for s in arm_samples:
                cmd_q = list(s["q"])
                for j in WAIST_INDICES:
                    cmd_q[j] = initial_q[j]            # keep waist rigid
                self.SendCommand(cmd_q, PREVIEW_ARM_KP, PREVIEW_ARM_KD, WAIST_KP, WAIST_KD)
                if self.record_hand and "h" in s:
                    self.SendHand(s["h"])
                time.sleep(1.0 / RECORD_HZ)
        except KeyboardInterrupt:
            print("\n⚠️  Preview interrupted — returning home.")

        print("✅ Preview done.")
        # 3) Return arms to home so it's safe to save / preview again
        self.ReturnToHome(initial_q, home_q, home_filename)

    def Run(self, filename: str, duration: float, home_filename: str):
        print("Waiting for robot connection...", end="", flush=True)
        try:
            while not self.first_state:
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("\n⛔ Ctrl+C pressed before connection was ready.")
            return False
        print(" Connected!")

        # Load Home Pose (same idea as replay)
        home_q = load_home_pose(home_filename)

        initial_q = [self.low_state.motor_state[i].q for i in range(G1_NUM_MOTOR)]

        print("\n" + "="*50)
        print("🟢 STABLE TEACH MODE (V4)")
        print(f"   - 0s - {SAFE_START_TIME}s:  ARMS LOCKED")
        print(f"   - Recording:    ARMS LIMP")
        print(f"   - Waist:        LOCKED (Rigid)")
        print(f"   - Return:       ARMS -> HOME ({home_filename})")
        print("="*50 + "\n")

        # Open the hand fully before recording starts (then state mode releases it).
        if self.record_hand:
            print("🖐️  Opening hand fully before recording...")
            for _ in range(40):
                self.SendHand([1.0] * HAND_MOTORS)
                time.sleep(0.02)
            self.hand_q = [1.0] * HAND_MOTORS

        # Push-to-teach: calibrate the resting force baseline (fingers open, untouched).
        if self.record_hand and self.hand_source == "push":
            print("🤛 Calibrating finger-force baseline — DON'T touch the fingers (~2s)...")
            self.hand_target = [1.0] * HAND_MOTORS
            acc = [0.0] * HAND_MOTORS
            fmin = [1e9] * HAND_MOTORS
            fmax = [-1e9] * HAND_MOTORS
            ncal = 0
            for _ in range(40):
                self.SendHand(self.hand_target, force=self.default_grip)
                if any(abs(v) > 1e-6 for v in self.hand_force):
                    for i in range(HAND_MOTORS):
                        acc[i] += self.hand_force[i]
                        fmin[i] = min(fmin[i], self.hand_force[i])
                        fmax[i] = max(fmax[i], self.hand_force[i])
                    ncal += 1
                time.sleep(0.05)
            if ncal:
                self.force_base = [acc[i] / ncal for i in range(HAND_MOTORS)]
                # Per-finger press threshold: honor the (low) user threshold, but never
                # below this finger's measured resting noise + margin, so a light real
                # press triggers while sensor noise doesn't cause phantom toggles.
                self.push_thresh_eff = [max(self.push_thresh, (fmax[i] - fmin[i]) * 1.5 + 20)
                                        for i in range(HAND_MOTORS)]
                print("   baseline(g):", [round(x) for x in self.force_base[:6]])
                print("   press-threshold R(g):", [round(x) for x in self.push_thresh_eff[:6]])
            else:
                self.push_thresh_eff = [self.push_thresh] * HAND_MOTORS
                print("   ⚠️  no force data (hand adapter dropped?) — push detection won't work this run.")
            self.pushed_prev = [False] * HAND_MOTORS

        t_start = time.time()
        dt = 1.0 / RECORD_HZ
        samples = []
        is_limp = False

        if self.record_hand:
            if self.hand_source == "push":
                if self.push_mode == "hold":
                    print("✋ Hand source = PUSH-to-teach (HOLD): press & HOLD a finger → it closes gradually; "
                          "release → holds at that level. Press again to open it back.")
                else:
                    print("✋ Hand source = PUSH-to-teach (TOGGLE): press a finger → flips full open/close.")
                print("   Keyboard: o/c = open/close BOTH hands · [ ] = open/close RIGHT · ; ' = open/close LEFT")
            elif self.hand_source == "follow":
                print(f"✋ Hand source = FOLLOW (force={self.follow_force:.0f}g): push the fingers by hand — "
                      "the motor gives way and follows; recording actual positions.")
                print("   (if a finger won't budge, raise/lower --follow-force, or use --hand-source keys)")
            elif self.hand_source == "state":
                print("✋ Hand source = STATE: legacy hold-last-cmd (fingers won't back-drive — prefer 'follow').")
            elif self.hand_source == "keys":
                print("✋ Hand source = KEYS:  o=open  c=close  b=box  t=thumbs-up")
            else:
                print("✋ Hand source = CMD: recording the rt/inspire/cmd stream (dashboard).")

        try:
            with KeyPoller() as kp:
                warned = False
                while True:
                    loop_start = time.time()
                    t_curr = loop_start - t_start
                    if t_curr > duration: break

                    if self.record_hand:
                        if self.hand_source == "push":
                            # Push-to-teach: the finger can't back-drive, but pressing it
                            # spikes its FORCE sensor. A press toggles that finger
                            # open<->close and the motor drives it there (recorded).
                            for i in range(HAND_MOTORS):
                                dev = self.hand_force[i] - self.force_base[i]
                                if abs(dev) > self.force_peak[i]:
                                    self.force_peak[i] = abs(dev)            # track peak for tuning
                                thr = self.push_thresh_eff[i]
                                pressing = abs(dev) > thr
                                if self.push_mode == "toggle":
                                    if pressing and not self.pushed_prev[i]:  # rising edge -> flip
                                        self.hand_target[i] = 0.0 if self.hand_target[i] > 0.5 else 1.0
                                        self.pushed_prev[i] = True
                                        print(f"\n  👉 {'close' if self.hand_target[i] < 0.5 else 'open'} "
                                              f"finger {i} (force {dev:+.0f}g)")
                                    elif abs(dev) < thr * 0.5:               # released (hysteresis)
                                        self.pushed_prev[i] = False
                                else:  # "hold" -> press-and-hold ramps the finger to any level
                                    if pressing:
                                        if not self.pushed_prev[i]:          # press start: pick direction
                                            self.push_dir[i] = -1.0 if self.hand_target[i] >= 0.5 else 1.0
                                            self.pushed_prev[i] = True
                                        speed = self.close_rate * min(3.0, abs(dev) / thr)  # harder = faster
                                        self.hand_target[i] = min(1.0, max(0.0,
                                            self.hand_target[i] + self.push_dir[i] * speed * dt))
                                    elif abs(dev) < thr * 0.5:               # released -> hold at this level
                                        if self.pushed_prev[i]:
                                            print(f"\n  ✋ finger {i} held at {self.hand_target[i]:.2f}")
                                        self.pushed_prev[i] = False
                            key = kp.poll()
                            if key:
                                kl = key.lower()
                                if kl == "o":                                   # open BOTH hands
                                    self.hand_target = [1.0] * HAND_MOTORS
                                    print("\n  ⌨️  open both hands")
                                elif kl == "c":                                 # close BOTH hands
                                    self.hand_target = [0.0] * HAND_MOTORS
                                    print("\n  ⌨️  close both hands")
                                elif kl == "[":                                 # open RIGHT hand
                                    for i in range(6): self.hand_target[i] = 1.0
                                    print("\n  ⌨️  open right hand")
                                elif kl == "]":                                 # close RIGHT hand
                                    for i in range(6): self.hand_target[i] = 0.0
                                    print("\n  ⌨️  close right hand")
                                elif kl == ";":                                 # open LEFT hand
                                    for i in range(6, 12): self.hand_target[i] = 1.0
                                    print("\n  ⌨️  open left hand")
                                elif kl == "'":                                 # close LEFT hand
                                    for i in range(6, 12): self.hand_target[i] = 0.0
                                    print("\n  ⌨️  close left hand")
                            self.SendHand(self.hand_target, force=self.default_grip)
                        elif self.hand_source == "follow":
                            # Low-force follow: command the CURRENT sensed position
                            # back at low grip force, so the motor tracks wherever you
                            # push each finger (no spring-back) and the actual motion
                            # is recorded. Needs the fingers to yield at this force.
                            self.SendHand(self.hand_state, force=self.follow_force)
                        elif self.hand_source == "state":
                            self.SendHand([-1.0] * HAND_MOTORS)  # (legacy) hold last cmd
                        elif self.hand_source == "keys":
                            key = kp.poll()
                            if key and key.lower() in HAND_PRESETS:
                                self.hand_q = list(HAND_PRESETS[key.lower()])
                                self.SendHand(self.hand_q)
                        # Warn if a hand isn't reporting (dropped adapter)
                        if self.hand_source in ("follow", "state", "push") and (not warned) and \
                           t_curr > SAFE_START_TIME + 1.0 and \
                           not any(abs(v) > 1e-3 for v in self.hand_state):
                            print("\n⚠️  Hand read-back all 0 (a hand's adapter is dropped) — that hand's track will be empty.")
                            print("    Reseat its USB-RS485 adapter (hub is flaky), or use --hand-source keys.")
                            warned = True

                    current_q_all = [float(self.low_state.motor_state[i].q) for i in range(G1_NUM_MOTOR)]
                    sample = {"t": round(t_curr, 4), "q": current_q_all}
                    if self.record_hand:
                        sample["h"] = self.hand_frame()
                    samples.append(sample)

                    if t_curr < SAFE_START_TIME:
                        print(f"\r💪 HOLDING... {SAFE_START_TIME - t_curr:.1f}s", end="")
                        self.SendCommand(initial_q, HOLD_ARM_KP, HOLD_ARM_KD, WAIST_KP, WAIST_KD)
                    else:
                        if not is_limp:
                            print("\n🔓 ARMS RELEASED! Move them now!")
                            is_limp = True

                        cmd_q = list(initial_q)
                        for i in ARM_INDICES: cmd_q[i] = 0.0
                        self.SendCommand(cmd_q, TEACH_ARM_KP, TEACH_ARM_KD, WAIST_KP, WAIST_KD)

                    time.sleep(max(0, dt - (time.time() - loop_start)))

        except KeyboardInterrupt:
            print("\n⛔ Recording stopped by Ctrl+C.")

        # Restore normal grip force after a low-force follow session (holds the
        # final recorded pose) so preview/replay and later use grip normally.
        if self.record_hand and self.hand_source == "follow":
            for _ in range(10):
                self.SendHand(self.hand_state, force=self.default_grip)
                time.sleep(0.02)

        # Push-to-teach tuning aid: how hard each finger registered this run. If a
        # finger you pressed didn't trigger, set --push-thresh below its peak.
        if self.record_hand and self.hand_source == "push":
            NM = ["R.pinky", "R.ring", "R.mid", "R.index", "R.thumbB", "R.thumbR",
                  "L.pinky", "L.ring", "L.mid", "L.index", "L.thumbB", "L.thumbR"]
            print(f"\n  peak force per finger (g)   [base threshold = {self.push_thresh:.0f}]:")
            for i in range(HAND_MOTORS):
                pk = self.force_peak[i]
                thr = self.push_thresh_eff[i]
                mark = "✓ triggered" if pk > thr else \
                       (f"↑ close (need >{thr:.0f}) — press the FINGERTIP PAD harder" if pk > 20 else "— no signal (dropped adapter?)")
                print(f"    {NM[i]:9s} {pk:5.0f}  (thr {thr:3.0f})  {mark}")

        self.ReturnToHome(initial_q, home_q, home_filename)

        # Offer to replay/preview the recording BEFORE saving (SDK still enabled).
        # Can be repeated; each preview returns the arms to home afterward.
        while samples:
            preview_choice = ask_yes_no("▶️  Replay this recording to preview before saving? [y/n]: ")
            if not preview_choice:   # False (no) or None (Ctrl+C)
                break
            self.Preview(samples, initial_q, home_q, home_filename)

        self.DisableSDK()

        total_time = samples[-1]["t"] if samples else 0.0
        save_choice = should_save_recording()
        if save_choice is None:
            print("🗑️ Recording not saved.")
            print(f"⏱️ Total recorded time: {total_time:.3f}s")
            print("✅ Done.")
            return False

        if not save_choice:
            print("🗑️ Recording not saved.")
            print(f"⏱️ Total recorded time: {total_time:.3f}s")
            repeat_choice = should_repeat_recording()
            if repeat_choice is None:
                print("✅ Done.")
                return False
            if repeat_choice:
                print("🔁 Restarting recording...")
                return True
            print("✅ Done.")
            return False

        with open(filename, "w") as f:
            meta = {"hz": RECORD_HZ, "motors": G1_NUM_MOTOR}
            if self.record_hand:
                meta["hand"] = True
                meta["hand_motors"] = HAND_MOTORS
            f.write(json.dumps({"meta": meta}) + "\n")
            for s in samples:
                f.write(json.dumps(s) + "\n")

        file_size_bytes = os.path.getsize(filename) if os.path.exists(filename) else 0
        file_size_mb = file_size_bytes / (1024 * 1024)

        print(f"💾 Saved {len(samples)} samples to {filename}")
        print(f"⏱️ Total recorded time: {total_time:.3f}s")
        print(f"📦 File size: {file_size_bytes} bytes ({file_size_mb:.3f} MB)")
        print("✅ Done.")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("iface", help="Network interface")
    parser.add_argument("--output", default="record", help="Output filename")
    parser.add_argument("--seconds", type=float, default=20.0)
    parser.add_argument("--home", default="arm_home.jsonl", help="Home pose file (jsonl)")
    parser.add_argument("--hand", dest="hand", action="store_true", default=True,
                        help="Also record the Inspire hand (default ON)")
    parser.add_argument("--no-hand", dest="hand", action="store_false", help="Do not record the hand")
    parser.add_argument("--hand-source", choices=["keys", "push", "cmd", "follow", "state"], default="keys",
                        help="keys=keyboard presets o/c/b/t during teach (default; reliable); "
                             "push=push-to-teach, press a finger to toggle it (uses force sensing); "
                             "cmd=dashboard rt/inspire/cmd stream; "
                             "follow=low-force by-hand (does NOT work on RH56 — non-backdrivable); "
                             "state=legacy hold-last-cmd")
    parser.add_argument("--follow-force", type=float, default=50.0,
                        help="grip force (grams) in follow mode; lower = easier to push by hand (default 50)")
    parser.add_argument("--push-thresh", type=float, default=50.0,
                        help="min force deviation (grams) for a press; auto-raised per finger to its resting noise (default 50)")
    parser.add_argument("--push-mode", choices=["hold", "toggle"], default="hold",
                        help="hold=press-and-hold ramps the finger to any level, release holds (default); "
                             "toggle=press flips full open/close")
    parser.add_argument("--close-rate", type=float, default=1.2,
                        help="finger close/open speed while held in push 'hold' mode (units/sec, default 1.2)")
    args = parser.parse_args()

    ChannelFactoryInitialize(0, args.iface)
    final_path = ensure_unique_path(args.output)
    teacher = StableTeacher()
    teacher.record_hand = args.hand
    teacher.hand_source = args.hand_source
    teacher.follow_force = args.follow_force
    teacher.push_thresh = args.push_thresh
    teacher.push_mode = args.push_mode
    teacher.close_rate = args.close_rate
    print(f"✋ Hand recording: {'ON (source=' + args.hand_source + ')' if args.hand else 'OFF'}")
    try:
        while teacher.Run(final_path, args.seconds, args.home):
            pass
    except KeyboardInterrupt:
        print("\n⛔ Ctrl+C pressed. Exiting cleanly.")
