#!/usr/bin/env python
"""Replay a LeRobot (egodex-style) bimanual trajectory on r1pro_dex in fastsim,
and record the result via the `record` extension.

The source dataset stores, per frame, a 56-D ``observation.state``:

    [ 0: 3]  left  wrist position (xyz)        |
    [ 3: 6]  left  wrist orientation (euler xyz)|  left  arm  -> solved with IK
    [ 6:28]  left  hand 22 joint angles         -> set directly
    [28:31]  right wrist position (xyz)         |
    [31:34]  right wrist orientation (euler xyz)|  right arm  -> solved with IK
    [34:56]  right hand 22 joint angles         -> set directly

Because the wrist poses are expressed in the capture/camera frame (unknown
transform to the robot base), the arms are driven by *delta poses*: each frame's
target EE pose in the robot base frame is

    target_base[t] = base_ee_init * ( cam_ee[0]^-1 * cam_ee[t] )

i.e. the trajectory's motion relative to its own first frame (expressed in the
EE-local frame, which is camera-frame-independent) applied on top of the robot's
actual initial EE pose. The 7-DoF arm joints are recovered with fastsim's
per-arm differential IK; the simulation runs physics-disabled (pure kinematic
playback) at dt=1/60 so each 15-fps data frame gets ~4 IK convergence steps.

Objects are NOT replayed (per requirement).

Usage (must use the fastsim conda env's python):
    python replay_lerobot.py --episode 0
    python replay_lerobot.py --all          # one subprocess per episode
"""
import argparse
import os
import subprocess
import sys
import tempfile

import numpy as np

# --------------------------------------------------------------------------- #
#  Constants
# --------------------------------------------------------------------------- #
HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_DATASET = os.path.join(HERE, "extracted", "add_remove_lid_15fps_10epi")
DEFAULT_CONFIG = os.path.join(HERE, "replay_config.yaml")
DATA_FPS = 15
ROBOT_NAME = "r1pro_dex"

# 22 hand joints per side, in the exact order the dataset packs them
# (matches r1pro_dex.usd / benchmark.yaml actuator lists: thumb5, index4,
#  middle4, ring4, pinky5).
LEFT_HAND_JOINTS = [
    "left_thumb_CMC_FE", "left_thumb_CMC_AA", "left_thumb_MCP_FE", "left_thumb_MCP_AA", "left_thumb_IP",
    "left_index_MCP_FE", "left_index_MCP_AA", "left_index_PIP", "left_index_DIP",
    "left_middle_MCP_FE", "left_middle_MCP_AA", "left_middle_PIP", "left_middle_DIP",
    "left_ring_MCP_FE", "left_ring_MCP_AA", "left_ring_PIP", "left_ring_DIP",
    "left_pinky_CMC", "left_pinky_MCP_FE", "left_pinky_MCP_AA", "left_pinky_PIP", "left_pinky_DIP",
]
RIGHT_HAND_JOINTS = [n.replace("left_", "right_", 1) for n in LEFT_HAND_JOINTS]

# 56-D observation.state slices
SLICES = {
    "left_arm":  {"pos": slice(0, 3),   "euler": slice(3, 6),   "hand": slice(6, 28)},
    "right_arm": {"pos": slice(28, 31), "euler": slice(31, 34), "hand": slice(34, 56)},
}
HAND_JOINTS = {"left_arm": LEFT_HAND_JOINTS, "right_arm": RIGHT_HAND_JOINTS}
ARMS = ["left_arm", "right_arm"]

# --- Arm targeting -------------------------------------------------------- #
# Delta-pose replay (camera pose is IGNORED): for each frame compute the relative
# transform from frame 0 in the data's own frame, then apply it on top of the
# robot's actual initial EE pose:
#     delta_t   = cam_ee[0]^-1 * cam_ee[t]      # relative motion, in EE-local frame
#     target_t  = ee_init * delta_t             # applied to the robot's init EE pose
# Frame 0 -> delta = identity -> robot starts exactly at its init EE pose.
# Optional fixed correction between the dataset wrist frame and r1pro's ee_link.
EE_OFFSET = {
    "left_arm":  np.eye(4),
    "right_arm": np.eye(4),
}


# --------------------------------------------------------------------------- #
#  Data loading
# --------------------------------------------------------------------------- #
def load_episode(dataset_dir: str, episode: int) -> np.ndarray:
    """Return (N, 56) float32 states for one episode, ordered by frame_index."""
    import pyarrow.parquet as pq

    data_file = os.path.join(dataset_dir, "data", "chunk-000", "file-000.parquet")
    df = pq.read_table(data_file).to_pandas()
    ep = df[df["episode_index"] == episode].sort_values("frame_index")
    if len(ep) == 0:
        raise ValueError(f"episode {episode} not found in {data_file}")
    states = np.stack(ep["observation.state"].to_numpy()).astype(np.float64)
    assert states.shape[1] == 56, f"expected 56-D state, got {states.shape}"
    return states


def build_temp_config(base_config: str, out_subdir: str) -> str:
    """Clone the YAML config, redirecting the recorder output to a per-episode dir."""
    import yaml

    with open(base_config) as f:
        cfg = yaml.safe_load(f)
    rec = cfg["extension"]["extension_cfg_dict"]["record"]
    rec["backend_root_path"] = f"output://{out_subdir}"
    fd, path = tempfile.mkstemp(suffix=".yaml", prefix="replay_cfg_")
    with os.fdopen(fd, "w") as f:
        yaml.safe_dump(cfg, f, sort_keys=False)
    return path


# --------------------------------------------------------------------------- #
#  Single-episode replay (runs inside fastsim)
# --------------------------------------------------------------------------- #
def run_episode(dataset_dir: str, episode: int, base_config: str):
    states = load_episode(dataset_dir, episode)
    n_frames = len(states)
    dataset_name = os.path.basename(dataset_dir.rstrip("/"))
    out_subdir = f"benchmark_replay_record/{dataset_name}/episode_{episode:02d}"
    temp_config = build_temp_config(base_config, out_subdir)

    # Heavy sim imports happen only now (after the lightweight data load).
    from fastsim.app import FastSim
    from fastsim.unisim.scene_manager import SceneManager
    from fastsim.utils.pose import Pose
    from fastsim.utils.log import Log

    # Pre-build camera-frame EE poses + hand targets per arm.
    cam_ee = {arm: [] for arm in ARMS}
    hand_targets = {arm: [] for arm in ARMS}
    for s in states:
        for arm in ARMS:
            sl = SLICES[arm]
            cam_ee[arm].append(
                Pose.from_euler_xyz(position=s[sl["pos"]].tolist(),
                                    euler_xyz=s[sl["euler"]].tolist())
            )
            hand_targets[arm].append(s[sl["hand"]].tolist())

    sim = FastSim(temp_config)
    sim.set_physics_disabled(True)        # pure kinematic playback

    ee_off = {arm: Pose.from_homogeneous_matrix(EE_OFFSET[arm]) for arm in ARMS}
    state = {"base_ee_init": {}, "cam0_inv": {}, "ready": False}

    def on_post_reset():
        """Anchor the delta-pose replay to the robot's actual init EE pose."""
        robot = SceneManager.get_robot(ROBOT_NAME)
        for arm in ARMS:
            state["base_ee_init"][arm] = robot.get_ee_pose(arm_name=arm)
            state["cam0_inv"][arm] = cam_ee[arm][0].inverse()
        state["ready"] = True
        Log.info(f"[replay] episode {episode}: {n_frames} frames @ {DATA_FPS}fps; "
                 f"delta-pose replay (camera pose ignored)", title="Replay")

    def apply_frame(robot, idx):
        for arm in ARMS:
            # delta = cam_ee[0]^-1 * cam_ee[idx];  target = ee_init * delta * ee_off
            delta = state["cam0_inv"][arm] * cam_ee[arm][idx]
            target = state["base_ee_init"][arm] * (delta * ee_off[arm])
            ik = robot.solve_ik(target, arm_name=arm)
            jpos = np.asarray(ik["joint_position"]).reshape(-1).tolist()
            robot.set_joint_position(jpos, ik["joint_names"])
            robot.set_joint_position_target(jpos, ik["joint_names"])
            # hand joints: exact angles, set directly
            robot.set_joint_position(hand_targets[arm][idx], HAND_JOINTS[arm])
            robot.set_joint_position_target(hand_targets[arm][idx], HAND_JOINTS[arm])

    def on_step(ctx):
        if not state["ready"]:
            return
        interval = max(1, round(1.0 / (ctx.dt * DATA_FPS)))
        # Allow one extra interval so the recorder (which captures every
        # `interval` steps) flushes the final converged frame.
        if ctx.step > (n_frames + 1) * interval:
            sim.request_terminate()
            return
        idx = min(ctx.step // interval, n_frames - 1)
        robot = SceneManager.get_robot(ROBOT_NAME)
        apply_frame(robot, idx)

    sim.add_post_reset_callback(on_post_reset)
    sim.add_step_callback(on_step)
    Log.info(f"[replay] output -> output://{out_subdir}", title="Replay")
    sim.start()   # setup + loop; calls os._exit(0) on finish


# --------------------------------------------------------------------------- #
#  Multi-episode driver (subprocess per episode, since FastSim os._exit()s)
# --------------------------------------------------------------------------- #
def run_all(dataset_dir: str, base_config: str, episodes):
    for ep in episodes:
        print(f"\n========== EPISODE {ep} ==========", flush=True)
        rc = subprocess.run(
            [sys.executable, os.path.abspath(__file__),
             "--dataset", dataset_dir, "--config", base_config, "--episode", str(ep)]
        ).returncode
        # FastSim exits via os._exit(0); a non-zero code means a real failure.
        if rc not in (0,):
            print(f"[replay] episode {ep} exited with code {rc}", flush=True)


def list_episodes(dataset_dir: str):
    import json
    info = json.load(open(os.path.join(dataset_dir, "meta", "info.json")))
    return list(range(int(info["total_episodes"])))


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--dataset", default=DEFAULT_DATASET, help="LeRobot dataset dir")
    ap.add_argument("--config", default=DEFAULT_CONFIG, help="fastsim YAML config")
    ap.add_argument("--episode", type=int, default=None, help="single episode index")
    ap.add_argument("--all", action="store_true", help="replay every episode")
    args = ap.parse_args()

    if args.all:
        run_all(args.dataset, args.config, list_episodes(args.dataset))
    else:
        ep = 0 if args.episode is None else args.episode
        run_episode(args.dataset, ep, args.config)


if __name__ == "__main__":
    main()