#!/usr/bin/env python

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
OpenArms End-Effector Replay Example with Visualization

Replays a dataset recorded with absolute joint positions by:
1. Converting joint positions to EE poses using FK
2. Converting EE poses back to joint positions using IK
3. Sending joint commands to the robot OR visualizing in simulation

Supports three modes:
- real: Send commands to physical robot
- sim: Visualize in simulation only (no robot required)
- both: Real robot + visualization

Example usage:
    python examples/openarms/replay_ee.py --mode sim
    python examples/openarms/replay_ee.py --mode real
    python examples/openarms/replay_ee.py --mode both --visualizer meshcat
"""

import argparse
import time
from os.path import dirname, expanduser

import numpy as np

from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
from lerobot.processor.converters import (
    robot_action_observation_to_transition,
    robot_action_to_transition,
    transition_to_robot_action,
)
from lerobot.robots.openarms.robot_kinematic_processor import (
    BimanualEEBoundsAndSafety,
    BimanualForwardKinematicsJointsToEE,
    BimanualInverseKinematicsEEToJoints,
)
from lerobot.utils.constants import ACTION
from lerobot.utils.robot_utils import precise_sleep


# Default configuration
DEFAULT_EPISODE_IDX = 0
DEFAULT_DATASET = "lerobot-data-collection/rac_blackf0"
DEFAULT_URDF = "src/lerobot/robots/openarms/urdf/openarm_bimanual_pybullet.urdf"
DEFAULT_LEFT_EE_FRAME = "openarm_left_hand_tcp"
DEFAULT_RIGHT_EE_FRAME = "openarm_right_hand_tcp"

# Motor names as used in the dataset actions (e.g., left_joint_1.pos)
MOTOR_NAMES = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "joint_7", "gripper"]

# URDF joint names (no underscore between "joint" and number)
LEFT_URDF_JOINTS = [f"openarm_left_joint{i}" for i in range(1, 8)]
RIGHT_URDF_JOINTS = [f"openarm_right_joint{i}" for i in range(1, 8)]


class MeshcatVisualizer:
    """Lightweight URDF visualizer using pinocchio + meshcat."""
    
    def __init__(self, urdf_path: str):
        import pinocchio as pin
        from pinocchio.visualize import MeshcatVisualizer as PinMeshcat
        
        urdf_dir = dirname(urdf_path)
        self.model, self.collision_model, self.visual_model = pin.buildModelsFromUrdf(
            urdf_path, urdf_dir, pin.JointModelFreeFlyer()
        )
        self.data = self.model.createData()
        
        self.viz = PinMeshcat(self.model, self.collision_model, self.visual_model)
        self.viz.initViewer(open=True)
        self.viz.loadViewerModel()
        
        # Build joint name mapping: dataset name -> pinocchio joint index
        # Dataset uses: left_joint_1, right_joint_2, etc.
        # URDF uses: openarm_left_joint1, openarm_right_joint2, etc.
        self.joint_map = {}
        for jid in range(1, self.model.njoints):
            urdf_name = self.model.names[jid]  # e.g., "openarm_left_joint1"
            # Extract side and number
            if "left_joint" in urdf_name:
                num = urdf_name.split("joint")[-1]  # "1"
                dataset_name = f"left_joint_{num}"
                self.joint_map[dataset_name] = jid
            elif "right_joint" in urdf_name:
                num = urdf_name.split("joint")[-1]
                dataset_name = f"right_joint_{num}"
                self.joint_map[dataset_name] = jid
        
        print(f"  Meshcat viewer opened (mapped {len(self.joint_map)} joints)")
        print(f"  Joint map: {list(self.joint_map.keys())[:4]}...")
        print("  Waiting for meshcat to load...")
        time.sleep(3)  # Give meshcat time to load meshes
        self._first_update = True
    
    def update(self, joint_positions: dict[str, float]):
        """Update visualization with new joint positions."""
        if self._first_update:
            pos_keys = [k for k in joint_positions.keys() if k.endswith(".pos")]
            print(f"  First update keys: {pos_keys[:4]}...")
            # Print sample values
            for k in pos_keys[:2]:
                print(f"    {k} = {joint_positions[k]:.2f}")
        
        # Build configuration vector (base pose + joints)
        # Free flyer base: [x, y, z, qx, qy, qz, qw]
        q = np.zeros(self.model.nq)
        q[3:7] = [0, 0, 0, 1]  # Identity quaternion
        
        matched = 0
        # Map joint positions using pre-built mapping
        for name, pos in joint_positions.items():
            if not name.endswith(".pos"):
                continue
            joint_name = name.removesuffix(".pos")  # e.g., "left_joint_1"
            
            jid = self.joint_map.get(joint_name)
            if jid is not None:
                idx = self.model.idx_qs[jid]
                if idx < len(q):
                    q[idx] = np.deg2rad(pos)
                    matched += 1
        
        if self._first_update:
            print(f"  Matched {matched} joints, q[7:14] = {q[7:14]}")
            self._first_update = False
        
        self.viz.display(q)


class RerunVisualizer:
    """Rerun-based visualizer for plots and EE trajectories."""
    
    def __init__(self, urdf_path: str = None, session_name: str = "openarms_replay"):
        import rerun as rr
        self.rr = rr
        rr.init(session_name)
        rr.spawn(memory_limit="10%")
        print("  Rerun viewer spawned (plots only, use --visualizer meshcat for 3D robot)")
    
    def update(self, joint_positions: dict[str, float], ee_poses: dict[str, float], frame_idx: int):
        """Log joint positions and EE poses."""
        self.rr.set_time("frame", sequence=frame_idx)
        
        # Log EE positions as colored spheres
        for prefix, color in [("left", [255, 100, 100]), ("right", [100, 100, 255])]:
            x, y, z = ee_poses.get(f"{prefix}_ee.x"), ee_poses.get(f"{prefix}_ee.y"), ee_poses.get(f"{prefix}_ee.z")
            if None not in (x, y, z):
                self.rr.log(f"ee/{prefix}", self.rr.Points3D([[x, y, z]], colors=[color], radii=[0.02]))
        
        # Log joint positions as time series
        for name, pos in joint_positions.items():
            if name.endswith(".pos"):
                self.rr.log(f"joints/{name}", self.rr.Scalars(pos))
        
        # Log EE poses as time series
        for name, val in ee_poses.items():
            self.rr.log(f"ee_plots/{name}", self.rr.Scalars(val))


def parse_args():
    parser = argparse.ArgumentParser(description="OpenArms EE Replay with Visualization")
    parser.add_argument("--mode", choices=["real", "sim", "both"], default="sim",
                        help="Execution mode: real (robot), sim (visualization), both")
    parser.add_argument("--visualizer", choices=["meshcat", "rerun", "none"], default="meshcat",
                        help="Visualization backend (meshcat shows 3D robot, rerun shows plots)")
    parser.add_argument("--dataset", type=str, default=DEFAULT_DATASET,
                        help="Dataset repo ID")
    parser.add_argument("--episode", type=int, default=DEFAULT_EPISODE_IDX,
                        help="Episode index to replay")
    parser.add_argument("--urdf", type=str, default=DEFAULT_URDF,
                        help="Path to URDF file")
    parser.add_argument("--left-ee-frame", type=str, default=DEFAULT_LEFT_EE_FRAME,
                        help="Left arm end-effector frame name in URDF")
    parser.add_argument("--right-ee-frame", type=str, default=DEFAULT_RIGHT_EE_FRAME,
                        help="Right arm end-effector frame name in URDF")
    parser.add_argument("--port-left", type=str, default="can2",
                        help="CAN port for left arm")
    parser.add_argument("--port-right", type=str, default="can3",
                        help="CAN port for right arm")
    parser.add_argument("--speed", type=float, default=1.0,
                        help="Playback speed multiplier")
    return parser.parse_args()


def main():
    args = parse_args()
    use_robot = args.mode in ["real", "both"]
    use_viz = args.mode in ["sim", "both"] and args.visualizer != "none"
    
    print("=" * 70)
    print("OpenArms EE Replay (FK -> IK Pipeline)")
    print("=" * 70)
    print(f"\nMode: {args.mode}")
    print(f"Visualizer: {args.visualizer}")
    print(f"Dataset: {args.dataset}")
    print(f"Episode: {args.episode}")
    print(f"Speed: {args.speed}x")
    print("=" * 70)

    robot = None
    viz = None
    
    # Resolve URDF path (handle relative and ~ paths)
    from pathlib import Path
    urdf_path = args.urdf
    if urdf_path.startswith("~"):
        urdf_path = expanduser(urdf_path)
    elif not Path(urdf_path).is_absolute():
        # Relative to workspace root
        urdf_path = str(Path(__file__).parent.parent.parent / urdf_path)
    
    # Initialize robot if needed
    if use_robot:
        from lerobot.robots.openarms.config_openarms_follower import OpenArmsFollowerConfig
        from lerobot.robots.openarms.openarms_follower import OpenArmsFollower
        
        print("\n[1/5] Initializing robot...")
        robot_config = OpenArmsFollowerConfig(
            port_left=args.port_left,
            port_right=args.port_right,
            can_interface="socketcan",
            id="openarms_follower",
            disable_torque_on_disconnect=True,
            max_relative_target=10.0,
        )
        robot = OpenArmsFollower(robot_config)
    else:
        print("\n[1/5] Skipping robot (sim mode)")

    # Initialize visualizer if needed
    if use_viz:
        print(f"\n[2/5] Initializing {args.visualizer} visualizer...")
        if args.visualizer == "meshcat":
            viz = MeshcatVisualizer(urdf_path)
        elif args.visualizer == "rerun":
            viz = RerunVisualizer(urdf_path)
    else:
        print("\n[2/5] Skipping visualization")

    # Initialize kinematics with URDF joint names
    print("\n[3/5] Initializing kinematics solvers...")
    
    left_kinematics = RobotKinematics(
        urdf_path=urdf_path,
        target_frame_name=args.left_ee_frame,
        joint_names=LEFT_URDF_JOINTS,
    )
    right_kinematics = RobotKinematics(
        urdf_path=urdf_path,
        target_frame_name=args.right_ee_frame,
        joint_names=RIGHT_URDF_JOINTS,
    )

    # Build pipelines - use motor names without gripper for the processor
    motor_names_no_gripper = [n for n in MOTOR_NAMES if n != "gripper"]
    
    joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
        steps=[
            BimanualForwardKinematicsJointsToEE(
                left_kinematics=left_kinematics,
                right_kinematics=right_kinematics,
                motor_names=MOTOR_NAMES,
            ),
        ],
        to_transition=robot_action_to_transition,
        to_output=transition_to_robot_action,
    )

    ee_to_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
        steps=[
            BimanualEEBoundsAndSafety(
                end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
                max_ee_step_m=0.10,
            ),
            BimanualInverseKinematicsEEToJoints(
                left_kinematics=left_kinematics,
                right_kinematics=right_kinematics,
                motor_names=MOTOR_NAMES,
                initial_guess_current_joints=False,
            ),
        ],
        to_transition=robot_action_observation_to_transition,
        to_output=transition_to_robot_action,
    )

    # Load dataset
    print(f"\n[4/5] Loading dataset '{args.dataset}'...")
    dataset = LeRobotDataset(args.dataset, episodes=[args.episode])
    episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == args.episode)
    
    if len(episode_frames) == 0:
        raise ValueError(f"No frames found for episode {args.episode}")
    
    print(f"  Found {len(episode_frames)} frames at {dataset.fps} FPS")
    
    action_features = dataset.features.get(ACTION, {})
    action_names = action_features.get("names", [])
    actions = episode_frames.select_columns(ACTION)

    # Connect robot if needed
    if use_robot:
        print("\n[5/5] Connecting to robot...")
        robot.connect(calibrate=False)
        if not robot.is_connected:
            raise RuntimeError("Robot failed to connect!")
    else:
        print("\n[5/5] Skipping robot connection (sim mode)")

    print("\n" + "=" * 70)
    print(f"Ready to replay! Mode: {args.mode}")
    print("=" * 70)
    
    if use_robot:
        input("\nPress ENTER to start...")
    else:
        print("\nStarting visualization playback...")
        time.sleep(1)

    # Simulated observation for sim-only mode
    sim_obs = {f"{prefix}_{motor}.pos": 0.0 
               for prefix in ["left", "right"] 
               for motor in MOTOR_NAMES}

    try:
        for idx in range(len(episode_frames)):
            loop_start = time.perf_counter()

            # Get observation
            if use_robot:
                robot_obs = robot.get_observation()
            else:
                robot_obs = sim_obs.copy()

            # Build joint action from dataset
            action_array = actions[idx][ACTION]
            joint_action = {}
            for i, name in enumerate(action_names):
                if name.endswith(".pos"):
                    joint_action[name] = float(action_array[i])

            # Convert: joints -> EE (FK)
            ee_action = joints_to_ee(joint_action.copy())

            # Convert: EE -> joints (IK)
            final_joint_action = ee_to_joints((ee_action.copy(), robot_obs))

            # Update simulated observation for next iteration
            if not use_robot:
                sim_obs.update(final_joint_action)

            # Send to robot
            if use_robot:
                robot.send_action(final_joint_action)

            # Update visualization with ORIGINAL dataset trajectory
            if viz:
                if isinstance(viz, MeshcatVisualizer):
                    viz.update(joint_action)  # Use original, not FK->IK reconstructed
                elif isinstance(viz, RerunVisualizer):
                    viz.update(joint_action, ee_action, idx)

            # Maintain replay rate
            loop_duration = time.perf_counter() - loop_start
            dt_s = (1.0 / dataset.fps / args.speed) - loop_duration
            if dt_s > 0:
                precise_sleep(dt_s)

            if (idx + 1) % 100 == 0:
                progress = (idx + 1) / len(episode_frames) * 100
                print(f"Progress: {idx + 1}/{len(episode_frames)} ({progress:.1f}%)")

        print(f"\n✓ Replayed {len(episode_frames)} frames")

    except KeyboardInterrupt:
        print("\n\nReplay interrupted")
    finally:
        if use_robot and robot:
            print("\nDisconnecting robot...")
            robot.disconnect()
        print("✓ Done!")


if __name__ == "__main__":
    main()