#!/usr/bin/env python3
"""
RM75 Robot Controller for URDF without actuators
Direct joint position control (kinematic mode)
"""

import mujoco
import mujoco.viewer
import numpy as np
import time
from pathlib import Path


class RM75Controller:
    def __init__(self, urdf_path: str ="/home/zl/Downloads/urdf_rm75/RM75-B.urdf", enable_viewer: bool = True):
        """
        Initialize RM75 robot simulation from URDF

        Args:
            urdf_path: Path to RM75-B.urdf file
            enable_viewer: Show visualization window
        """
        # Load model
        self.model = mujoco.MjModel.from_xml_path(urdf_path)
        self.data = mujoco.MjData(self.model)

        # Robot info
        self.n_joints = self.model.njnt
        self.n_actuators = self.model.nu

        print(f"✓ Loaded RM75 robot")
        print(f"  - Joints: {self.n_joints}")
        print(f"  - Actuators: {self.n_actuators} (using direct joint control)")
        print(f"  - Bodies: {self.model.nbody}")

        # Get joint names for reference
        self.joint_names = []
        for i in range(self.n_joints):
            self.joint_names.append(self.model.joint(i).name)
        print(f"  - Joints: {', '.join(self.joint_names)}")

        # Store home position (current joint angles)
        self.home_position = self.data.qpos[:self.n_joints].copy()
        print(f"  - Home position: {self.home_position}")

        # For position control without actuators, we'll use qpos directly
        self.use_actuators = self.n_actuators > 0

        # Viewer
        self.viewer = None
        if enable_viewer:
            try:
                self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
                print("✓ Viewer launched successfully")
            except Exception as e:
                print(f"Warning: Could not launch viewer: {e}")
                self.viewer = None

    def get_joint_positions(self):
        """Get current joint angles (radians)"""
        return self.data.qpos[:self.n_joints].copy()

    def get_joint_velocities(self):
        """Get current joint velocities (rad/s)"""
        return self.data.qvel[:self.n_joints].copy()

    def get_end_effector_pose(self):
        """Get end-effector position and orientation"""
        # Last body is usually end-effector
        end_effector_id = self.model.nbody - 1
        position = self.data.xpos[end_effector_id].copy()
        orientation = self.data.xmat[end_effector_id].copy().reshape(3, 3)
        return position, orientation

    def set_joint_positions(self, positions):
        """
        Set joint positions directly (kinematic control)

        Args:
            positions: Target joint angles in radians (length should match n_joints)
        """
        if len(positions) != self.n_joints:
            print(f"Warning: Expected {self.n_joints} joints, got {len(positions)}")
            positions = positions[:self.n_joints]

        # Directly set joint positions (kinematic control)
        self.data.qpos[:self.n_joints] = positions

        # Also set velocities to zero to avoid unwanted motion
        self.data.qvel[:self.n_joints] = 0

    def step(self):
        """Advance simulation one step"""
        mujoco.mj_step(self.model, self.data)
        if self.viewer:
            self.viewer.sync()

    def step_n(self, n_steps: int):
        """Advance simulation by N steps"""
        for _ in range(n_steps):
            self.step()

    def move_to_position(self, target_positions, steps=500):
        """
        Smoothly move to target joint positions

        Args:
            target_positions: Target joint angles
            steps: Number of simulation steps for the movement
        """
        current = self.get_joint_positions()
        target = np.array(target_positions[:self.n_joints])

        for i in range(steps):
            # Linear interpolation
            alpha = (i + 1) / steps
            positions = current + alpha * (target - current)
            self.set_joint_positions(positions)
            self.step()

        # Ensure exact target
        self.set_joint_positions(target)
        self.step_n(10)

    def run_trajectory(self, trajectory_points, steps_between_points=500):
        """
        Execute a joint trajectory

        Args:
            trajectory_points: List of joint position arrays
            steps_between_points: Steps between each point
        """
        print(f"Executing trajectory with {len(trajectory_points)} points...")

        for i, target in enumerate(trajectory_points):
            print(f"  Moving to point {i + 1}/{len(trajectory_points)}")
            self.move_to_position(target, steps_between_points)

        print("✓ Trajectory complete")

    def run_forever(self, dt=0.01):
        """Run simulation with real-time control loop"""
        print("\n✓ Simulation running. Close viewer window to exit.\n")

        try:
            while self.viewer and self.viewer.is_running():
                self.step()
                time.sleep(dt)
        except KeyboardInterrupt:
            print("\n✓ Stopped by user")
        finally:
            if self.viewer:
                self.viewer.close()

    def print_state(self):
        """Print current robot state"""
        positions = self.get_joint_positions()
        print(f"Joint positions (rad): {[f'{p:.3f}' for p in positions]}")


# Test trajectories for RM75
def create_sine_wave_trajectory(controller, duration_seconds=5, frequency=0.5):
    """Create a sine wave trajectory for testing"""
    steps = int(duration_seconds / controller.model.opt.timestep)
    trajectory = []

    for i in range(steps):
        t = i * controller.model.opt.timestep
        positions = controller.home_position.copy()

        # Create sine wave motion on first 3 joints (shoulder, elbow, wrist)
        positions[0] = controller.home_position[0] + 0.5 * np.sin(2 * np.pi * frequency * t)
        positions[1] = controller.home_position[1] + 0.3 * np.sin(2 * np.pi * frequency * t + 1.0)
        positions[2] = controller.home_position[2] + 0.2 * np.sin(2 * np.pi * frequency * t + 2.0)

        trajectory.append(positions)

    return trajectory


if __name__ == "__main__":
    # Path to your URDF
    urdf_file = "/home/zl/Downloads/urdf_rm75/RM75-B.urdf"

    # Check if file exists
    if not Path(urdf_file).exists():
        print(f"Error: URDF file not found at {urdf_file}")
        exit(1)

    # Create robot controller
    print("Initializing RM75 Controller...")
    robot = RM75Controller(urdf_file, enable_viewer=True)

    if robot.viewer is None:
        print("Error: Could not initialize viewer. Running without visualization.")

    # Give time for viewer to initialize
    time.sleep(1)

    # Test 1: Print current state
    print("\n>>> Current robot state:")
    robot.print_state()

    # Test 2: Move to home position
    print("\n>>> Moving to home position...")
    robot.move_to_position(robot.home_position, steps=300)
    robot.print_state()

    # Test 3: Create and execute a pose sequence
    print("\n>>> Testing different poses...")

    # Pose 1: Slightly raised arm
    pose1 = robot.home_position.copy()
    pose1[0] = 0.5  # Joint 1
    pose1[1] = -0.3  # Joint 2
    pose1[2] = 0.2  # Joint 3

    # Pose 2: Extended arm
    pose2 = robot.home_position.copy()
    pose2[0] = 0.8
    pose2[1] = -0.5
    pose2[2] = 0.4
    pose2[3] = 0.3  # Joint 4

    # Pose 3: Folded position
    pose3 = robot.home_position.copy()
    pose3[0] = -0.5
    pose3[1] = 0.3
    pose3[2] = -0.2

    trajectory = [pose1, pose2, pose3, robot.home_position]
    robot.run_trajectory(trajectory, steps_between_points=400)

    # Test 4: Continuous sine wave motion (5 seconds)
    print("\n>>> Starting sine wave motion (5 seconds)...")
    sine_trajectory = create_sine_wave_trajectory(robot, duration_seconds=5, frequency=0.8)
    robot.run_trajectory(sine_trajectory, steps_between_points=1)

    # Return to home
    print("\n>>> Returning to home position...")
    robot.move_to_position(robot.home_position, steps=300)

    print("\n✓ Demo complete!")

    # Keep viewer open until user closes
    if robot.viewer:
        print("\nPress Ctrl+C in terminal to exit, or close the viewer window.")
        robot.run_forever()
    else:
        print("\nSimulation completed.")