#!/usr/bin/env python3
"""
RM75 Robot Controller with True Dynamics for URDF without Actuators
Run this in your coppeliasim conda environment
"""

import mujoco
import mujoco.viewer
import numpy as np
import time
from pathlib import Path
from scipy.signal import butter, lfilter


class RM75Controller:
    def __init__(self, urdf_path: str, enable_viewer: bool = True,
                 control_mode='torque',  # 'torque' or 'position'
                 low_pass_cutoff_hz=10.0):
        """
        Initialize RM75 robot simulation with dynamic control

        Args:
            urdf_path: Path to RM75-B.urdf file
            enable_viewer: Show visualization window
            control_mode: 'torque' (realistic dynamics) or 'position' (kinematic)
            low_pass_cutoff_hz: Cutoff frequency for command filtering (Hz)
        """
        # 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 in URDF: {self.n_actuators}")
        print(f"  - Bodies: {self.model.nbody}")

        # Get joint names and limits
        self.joint_names = []
        self.joint_lower_limits = []
        self.joint_upper_limits = []

        for i in range(self.n_joints):
            self.joint_names.append(self.model.joint(i).name)
            # Get joint limits from URDF (range is [lower, upper])
            self.joint_lower_limits.append(self.model.jnt_range[i, 0])
            self.joint_upper_limits.append(self.model.jnt_range[i, 1])

        print(f"  - Joint limits: {self.joint_lower_limits[0]:.2f} to {self.joint_upper_limits[0]:.2f} rad for joint_1")

        # If no actuators, we need to add them programmatically
        if self.n_actuators == 0:
            print("  No actuators in URDF. Adding torque actuators for dynamic control...")
            self._add_actuators_programmatically()

        # Control parameters
        self.control_mode = control_mode

        # PD gains for torque control (tuned for realistic motion)
        self.kp = np.array([200.0, 200.0, 150.0, 100.0, 80.0, 60.0, 50.0])  # Proportional gains
        self.kd = np.array([20.0, 20.0, 15.0, 10.0, 8.0, 6.0, 5.0])  # Derivative gains

        # State variables
        self.desired_positions = self.get_joint_positions().copy()

        # Low-pass filter for smooth commands
        self.lp_cutoff = low_pass_cutoff_hz
        self.filtered_commands = self.get_joint_positions().copy()

        # Home position
        self.home_position = self.get_joint_positions().copy()
        print(f"  - Home position: {self.home_position}")

        # Setup 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 _add_actuators_programmatically(self):
        """Add torque actuators programmatically to enable dynamic control"""
        # Save current model state
        n_new_actuators = self.n_joints

        # We need to create a new model with actuators
        # For now, we'll use direct qpos control with custom dynamics
        # This is a workaround by setting control_mode to 'position' for kinematic control
        # and using manual velocity/acceleration limits

        print(f"  Adding {n_new_actuators} virtual torque actuators")

        # Alternative approach: Use qfrc_applied to apply forces directly
        # This bypasses the need for actuators in the URDF
        self.use_qfrc_applied = True
        self.n_actuators = n_new_actuators

        # Create virtual control array
        self.virtual_ctrl = np.zeros(self.n_joints)

    def _low_pass_filter(self, new_target):
        """Apply simple first-order low-pass filter to target positions"""
        dt = self.model.opt.timestep
        alpha = 2 * np.pi * self.lp_cutoff * dt
        alpha = min(alpha, 1.0)  # Clamp for stability

        for i in range(self.n_joints):
            self.filtered_commands[i] = alpha * new_target[i] + (1 - alpha) * self.filtered_commands[i]

        return self.filtered_commands.copy()

    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 set_desired_positions(self, target_positions, apply_filter=True):
        """
        Set desired joint positions

        Args:
            target_positions: Target joint angles in radians
            apply_filter: Apply low-pass filter for smooth motion
        """
        target = np.array(target_positions[:self.n_joints])

        # Apply joint limits
        for i in range(self.n_joints):
            target[i] = np.clip(target[i], self.joint_lower_limits[i], self.joint_upper_limits[i])

        # Apply low-pass filter if enabled
        if apply_filter:
            target = self._low_pass_filter(target)

        self.desired_positions = target

    def compute_torques(self):
        """Compute torques using PD control with velocity damping"""
        current_pos = self.get_joint_positions()
        current_vel = self.get_joint_velocities()

        # Position error
        pos_error = self.desired_positions - current_pos

        # PD control law (negative feedback)
        torques = self.kp[:self.n_joints] * pos_error - self.kd[:self.n_joints] * current_vel

        # Apply torque limits (safety)
        max_torque = 50.0  # Nm limit
        torques = np.clip(torques, -max_torque, max_torque)

        return torques

    def step(self):
        """Step the simulation with control applied"""
        if self.control_mode == 'torque':
            # Compute torques
            torques = self.compute_torques()

            # Apply torques directly to joints using qfrc_applied
            # This bypasses the need for actuators in the URDF
            self.data.qfrc_applied[:self.n_joints] = torques

        elif self.control_mode == 'position':
            # Direct position control (kinematic)
            # Add velocity damping for smoother motion
            current_pos = self.get_joint_positions()
            pos_error = self.desired_positions - current_pos

            # Simple proportional velocity control
            kp_vel = 50.0
            target_vel = kp_vel * pos_error

            # Limit velocity
            max_vel = 3.0  # rad/s
            target_vel = np.clip(target_vel, -max_vel, max_vel)

            # Apply velocity
            self.data.qvel[:self.n_joints] = target_vel

            # For position mode, we also apply small corrective torques
            kp_correct = 100.0
            kd_correct = 10.0
            correction = kp_correct * pos_error - kd_correct * self.get_joint_velocities()
            self.data.qfrc_applied[:self.n_joints] = correction

        # Step physics
        mujoco.mj_step(self.model, self.data)

        # If using direct qpos control (position mode extreme)
        if self.control_mode == 'position_direct':
            self.data.qpos[:self.n_joints] = self.desired_positions
            self.data.qvel[:self.n_joints] = 0
            mujoco.mj_step(self.model, self.data)

        # Sync viewer if active
        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, duration=1.0, apply_filter=True):
        """
        Smoothly move to target position with dynamics

        Args:
            target_positions: Target joint angles
            duration: Movement duration (seconds)
            apply_filter: Apply low-pass filtering
        """
        # Calculate number of steps for smooth interpolation
        n_steps = int(duration / self.model.opt.timestep)
        start_positions = self.get_joint_positions()
        target = np.array(target_positions[:self.n_joints])

        print(f"  Moving with dynamics: {n_steps} steps over {duration}s")

        for i in range(n_steps):
            # Linear interpolation for desired positions
            alpha = (i + 1) / n_steps
            desired = start_positions + alpha * (target - start_positions)

            # Set desired position (filter will apply smoothing)
            self.set_desired_positions(desired, apply_filter=apply_filter)

            # Step simulation
            self.step()

            # Optional: print progress
            if (i + 1) % 200 == 0:
                progress = (i + 1) / n_steps * 100
                print(f"    Progress: {progress:.0f}%")

    def run_trajectory(self, trajectory_points, duration_per_point=1.0, apply_filter=True):
        """
        Execute a trajectory with smooth dynamics

        Args:
            trajectory_points: List of joint position arrays
            duration_per_point: Time per trajectory point (seconds)
            apply_filter: Apply low-pass filtering
        """
        print(f"\nExecuting 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, duration_per_point, apply_filter)

        print("✓ Trajectory complete")

    def run_interactive(self):
        """Run simulation with real-time control"""
        print("\n✓ Simulation running with dynamic control")
        print(f"  Control mode: {self.control_mode}")
        print(f"  Low-pass cutoff: {self.lp_cutoff} Hz")
        print("  Close viewer window to exit.\n")

        try:
            last_time = time.time()
            while self.viewer and self.viewer.is_running():
                self.step()

                # Maintain real-time speed
                elapsed = time.time() - last_time
                sleep_time = self.model.opt.timestep - elapsed
                if sleep_time > 0:
                    time.sleep(sleep_time)
                last_time = time.time()

        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()
        velocities = self.get_joint_velocities()
        print(f"Positions (rad): {[f'{p:.3f}' for p in positions[:4]]}...")
        print(f"Velocities (rad/s): {[f'{v:.3f}' for v in velocities[:4]]}...")


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

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

    print("=" * 60)
    print("RM75 Robot with Realistic Dynamics")
    print("=" * 60)

    # Test 1: Torque control with low-pass filter (most realistic)
    print("\n>>> Test 1: Torque Control with Low-Pass Filter (5 Hz)")
    robot1 = RM75Controller(urdf_file, enable_viewer=True,
                            control_mode='torque',
                            low_pass_cutoff_hz=5.0)  # Smooth, natural response

    # Wait a moment for viewer to initialize
    time.sleep(1)

    # Move to a pose with dynamics
    target_pose = robot1.home_position.copy()
    target_pose[0] = 0.8  # Joint 1 (base rotation)
    target_pose[1] = -0.5  # Joint 2 (shoulder)
    target_pose[2] = 0.4  # Joint 3 (elbow)
    target_pose[3] = 0.3  # Joint 4 (wrist 1)

    print("\nMoving to target pose with realistic dynamics...")
    robot1.move_to_position(target_pose, duration=2.0, apply_filter=True)

    print("\nRobot state after movement:")
    robot1.print_state()

    # Test 2: Return to home with different filter
    print("\n>>> Returning to home with softer filter (2 Hz)")
    robot1.lp_cutoff = 2.0  # Change to slower response
    robot1.move_to_position(robot1.home_position, duration=2.0, apply_filter=True)

    # Test 3: Demonstrate different control modes
    print("\n>>> Testing position control mode (faster response)")
    robot2 = RM75Controller(urdf_file, enable_viewer=False,
                            control_mode='position',
                            low_pass_cutoff_hz=15.0)

    test_pose = robot2.home_position.copy()
    test_pose[0] = 0.5
    test_pose[1] = -0.3

    print("Moving in position control mode...")
    robot2.move_to_position(test_pose, duration=1.0, apply_filter=True)
    robot2.print_state()

    print("\n" + "=" * 60)
    print("Dynamic Behavior Summary")
    print("=" * 60)
    print("✓ Torque Control: Realistic physics with inertia and damping")
    print("✓ Low-pass filter: Adjustable smoothing (lower = smoother)")
    print("✓ Joint limits: Automatically enforced from URDF")
    print("✓ Programmatic actuation: No URDF modification needed")
    print("\nTips for tuning:")
    print("  - Lower cutoff (2-5 Hz): Smooth, natural motion")
    print("  - Higher cutoff (10-20 Hz): More responsive")
    print("  - Adjust kp/kd gains for stiffness/damping")

    # Run interactive simulation
    print("\nStarting interactive simulation...")
    print("Press Ctrl+C in terminal to exit, or close the viewer window.")
    robot1.run_interactive()