#!/usr/bin/env python3
"""
G1 REPLAY V14 (HOME RETURN + FULL BODY LOCK)
--------------------------------------------
1. LOCK: Legs/Waist Rigid (Kp=300) to support weight.
2. REPLAY: Plays your recorded motion.
3. RETURN: Smoothly moves arms to 'arm_home.jsonl' pose at the end.
   - Ignores leg data from home file (keeps them locked for balance).
   - Slow 3-second transition for safety.

Usage:
  python3 g1_replay_v4_stable.py enp3s0 --input my_teleop_data.jsonl --home arm_home.jsonl
  
  
  python3 g1_replay_v4_stable.py enp3s0 --input change_battery.jsonl 
"""

import time
import sys
import json
import argparse
import numpy as np
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 IDL
from unitree_sdk2py.idl.unitree_go.msg.dds_ import MotorCmds_
from unitree_sdk2py.idl.default import unitree_go_msg_dds__MotorCmd_

# ✅ added (for pause/resume keys)
import termios
import tty
import select

G1_NUM_MOTOR = 29
ENABLE_ARM_SDK_INDEX = 29
DATA_DIR = Path("DataG1")
REPLAY_HZ = 60.0
HAND_TOPIC = "rt/inspire/cmd"
HAND_MOTORS = 12

# --- GAINS (ROBOT_ARM.PY STANDARD) ---
KP_HIGH = 300.0  # Core/Legs
KD_HIGH = 3.0
KP_LOW  = 80.0   # Arms/Ankles
KD_LOW  = 3.0
KP_WRIST = 40.0
KD_WRIST = 1.5

WEAK_MOTORS = [4, 10, 15, 16, 17, 18, 22, 23, 24, 25]
WRIST_MOTORS = [19, 20, 21, 26, 27, 28]

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
    print("⚠️  Using Default Home (Arms at 0.0)")
    default_q = [0.0] * G1_NUM_MOTOR
    return default_q

# ✅ added: non-blocking key reader (x pause, z play)
class KeyPoller:
    def __enter__(self):
        self.fd = sys.stdin.fileno()
        self.old_settings = termios.tcgetattr(self.fd)
        tty.setcbreak(self.fd)  # immediate key reads
        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

class ReplayWithHome:
    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 command publisher (drives the Inspire hand alongside the arm).
        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)]
        self.has_hand = False

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

    def SendHand(self, h):
        if not h:
            return
        for i in range(min(HAND_MOTORS, len(h))):
            self.hand_msg.cmds[i].q = float(h[i])
        self.hand_pub.Write(self.hand_msg)

    def SendFrame(self, arm_target_q, body_lock_q):
        self.low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 1.0
        
        for i in range(G1_NUM_MOTOR):
            self.low_cmd.motor_cmd[i].mode = 1
            self.low_cmd.motor_cmd[i].dq = 0
            self.low_cmd.motor_cmd[i].tau = 0
            
            # --- POSITIONS ---
            # Arms (15-28) -> Follow Target (Replay or Home)
            # Body (0-14)  -> Follow Lock (Statue Mode)
            if i >= 15: 
                self.low_cmd.motor_cmd[i].q = arm_target_q[i]
            else:       
                self.low_cmd.motor_cmd[i].q = body_lock_q[i]

            # --- GAINS ---
            if i in WEAK_MOTORS:
                self.low_cmd.motor_cmd[i].kp = KP_LOW
                self.low_cmd.motor_cmd[i].kd = KD_LOW
            elif i in WRIST_MOTORS:
                self.low_cmd.motor_cmd[i].kp = KP_WRIST
                self.low_cmd.motor_cmd[i].kd = KD_WRIST
            else:
                self.low_cmd.motor_cmd[i].kp = KP_HIGH # 300.0
                self.low_cmd.motor_cmd[i].kd = KD_HIGH

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

    def DisableSDK(self):
        print("\n🔌 Disabling SDK...")
        self.low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 0.0
        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        for _ in range(10):
            self.arm_pub.Write(self.low_cmd)
            time.sleep(0.02)

    def Run(self, filename: str, home_filename: str, speed: float):
        print("Waiting for robot...", end="", flush=True)
        while not self.first_state: time.sleep(0.1)
        print(" Connected!")

        # 1. LOAD DATA
        home_q = load_home_pose(home_filename)
        full_body_lock_q = [self.low_state.motor_state[i].q for i in range(G1_NUM_MOTOR)]
        
        frames = []
        try:
            with open(filename, 'r') as f:
                for line in f:
                    d = json.loads(line)
                    if 'q' in d: frames.append(d)
        except Exception as e: print(f"Error: {e}"); return
        
        self.has_hand = any('h' in fr for fr in frames)
        print(f"🟢 Ready to play {len(frames)} frames.")
        if self.has_hand:
            print("✋ Hand track found — driving rt/inspire/cmd alongside the arm.")
        print(f"🔒 Body is LOCKED. Arms will return to 'arm_home' at end.")
        print("🎮 Controls:  x = PAUSE   z = PLAY/RESUME   q = QUIT")
        input("👉 Press Enter to Begin...")

        # 2. MOVE TO START
        print("Moving to start...")
        file_start_q = frames[0]['q']
        file_start_h = frames[0].get('h')
        steps = 60
        for k in range(steps):
            alpha = k / steps
            interp_q = list(full_body_lock_q)
            for j in range(15, 29):
                interp_q[j] = (1-alpha)*full_body_lock_q[j] + alpha*file_start_q[j]
            self.SendFrame(interp_q, full_body_lock_q)
            if self.has_hand: self.SendHand(file_start_h)
            time.sleep(1.0/REPLAY_HZ)

        # 3. PLAY REPLAY (with pause/resume)
        print("▶️ Playing... (press x to pause)")
        last_played_q = file_start_q
        last_played_h = file_start_h
        paused = False

        # Instead of using (time.time() - t_start) directly, we use a play clock
        # that only advances when not paused.
        play_elapsed = 0.0
        last_real = time.time()

        try:
            with KeyPoller() as kp:
                while True:
                    # --- key handling ---
                    key = kp.poll()
                    if key:
                        key = key.lower()
                        if key == 'x' and not paused:
                            paused = True
                            print("\n⏸️  PAUSED (press z to continue)")
                        elif key == 'z' and paused:
                            paused = False
                            last_real = time.time()  # reset real-time anchor
                            print("\n▶️ RESUMED")
                        elif key == 'q':
                            print("\n🛑 QUIT")
                            break

                    now_real = time.time()

                    if paused:
                        # Hold last pose, do not advance play time
                        self.SendFrame(last_played_q, full_body_lock_q)
                        if self.has_hand: self.SendHand(last_played_h)
                        time.sleep(1.0 / REPLAY_HZ)
                        continue

                    # advance play clock only when running
                    dt_real = now_real - last_real
                    last_real = now_real
                    play_elapsed += dt_real * speed

                    target_frame = None
                    for f in frames:
                        if f['t'] - frames[0]['t'] >= play_elapsed:
                            target_frame = f
                            break
                    if target_frame is None:
                        break

                    self.SendFrame(target_frame['q'], full_body_lock_q)
                    last_played_q = target_frame['q']
                    if self.has_hand:
                        last_played_h = target_frame.get('h', last_played_h)
                        self.SendHand(last_played_h)
                    time.sleep(1.0 / REPLAY_HZ)

        except KeyboardInterrupt:
            print("\nStopped.")

        # 4. RETURN TO ARM HOME (Slow & Smooth)
        print(f"\n🏡 Returning arms to {home_filename}...")
        
        # 3 Seconds duration for smoothness
        home_steps = 180 
        
        for k in range(home_steps):
            alpha = k / home_steps
            # Interpolate: Last Pose -> Home Pose
            interp_q = list(last_played_q)
            for j in range(15, 29):
                interp_q[j] = (1-alpha)*last_played_q[j] + alpha*home_q[j]
            
            # Send (Body still locked to original standing pose)
            self.SendFrame(interp_q, full_body_lock_q)
            if self.has_hand: self.SendHand(last_played_h)
            time.sleep(1.0/REPLAY_HZ)

        print("✅ Home Reached.")
        self.DisableSDK()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("iface", help="Network interface")
    parser.add_argument("--input", required=True, help="Input recording file")
    parser.add_argument("--home", default="arm_home.jsonl", help="Home pose file")
    parser.add_argument("--speed", type=float, default=1.0)
    args = parser.parse_args()

    ChannelFactoryInitialize(0, args.iface)
    path = resolve_input_path(args.input)
    ReplayWithHome().Run(path, args.home, args.speed)
