#!/usr/bin/env python3
"""
Pure Position Control for MuJoCo - No velocity commands, no forces
Direct joint position control with smoothing
"""

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


class MuJoCoPositionController:
    """
    Pure position control - directly sets joint positions
    No velocity commands, no forces - completely stable
    """

    def __init__(self, urdf_path="./urdf_rm75/RM75-B.urdf", smoothness=0.05, enable_viewer=True):
        """
        Args:
            urdf_path: Path to URDF file
            smoothness: Motion smoothness (0.02=very smooth, 0.1=fast)
            enable_viewer: Show MuJoCo viewer
        """
        # Load model
        self.model = mujoco.MjModel.from_xml_path(urdf_path)
        self.data = mujoco.MjData(self.model)

        self.time_interval = 0.02

        print(f'time interval: {self.model.opt.timestep}')

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

        # Get joint limits
        self.joint_lower_limits = []
        self.joint_upper_limits = []
        for i in range(self.n_joints):
            self.joint_lower_limits.append(self.model.jnt_range[i, 0])
            self.joint_upper_limits.append(self.model.jnt_range[i, 1])

        print(f"Loaded robot: {self.n_joints} joints")
        for i in range(self.n_joints):
            print(
                f"  {self.model.joint(i).name}: limit [{self.joint_lower_limits[i]:.2f}, {self.joint_upper_limits[i]:.2f}]")

        # Target joint angles (in radians)
        self.target_joints = self.data.qpos[:self.n_joints].copy()

        # Smoothing factor (0-1, lower = smoother)
        self.smoothness = smoothness

        # Thread safety
        self.command_lock = threading.Lock()
        self.feedback_lock = threading.Lock()
        self.current_feedback_joint = self.data.qpos[:self.n_joints].copy()

        self.max_ang_inc = 0.02

        # Control flags
        self.running = False
        self.simulation_thread = None

        # Viewer
        self.viewer = None
        if enable_viewer:
            try:
                self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
                print("Viewer launched")
            except Exception as e:
                print(f"Viewer warning: {e}")
        self.start()

    def start(self):
        """Start the simulation thread"""
        if self.running:
            return

        self.running = True
        self.simulation_thread = threading.Thread(target=self._simulation_loop, daemon=True)
        self.simulation_thread.start()
        print("Simulation thread started")

    def stop(self):
        """Stop the simulation thread"""
        self.running = False
        if self.simulation_thread:
            self.simulation_thread.join(timeout=2.0)
        if self.viewer:
            self.viewer.close()
        print("Simulation stopped")

    def send_command(self, joint_positions):
        """
        Send target joint positions

        Args:
            joint_positions: Array of target joint angles (radians)
        """
        cmd = np.array(joint_positions[:self.n_joints], dtype=np.float64)

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

        with self.command_lock:
            self.target_joints = cmd

    def get_feedback(self):
        """Get current joint positions"""
        with self.feedback_lock:
            return self.current_feedback_joint.copy()

    def get_target(self):
        """Get current target positions"""
        with self.command_lock:
            return self.target_joints.copy()

    def _simulation_loop(self):
        """
        Main simulation loop - PURE POSITION CONTROL
        No velocity commands, no forces - just direct position setting
        """
        last_time = time.time()

        # For smooth interpolation
        current_joints = self.data.qpos[:self.n_joints].copy()

        while self.running:
            # Get target command
            with self.command_lock:
                target = self.target_joints.copy()

            # Get current positions
            current_joints = self.data.qpos[:self.n_joints].copy()

            # Smooth interpolation toward target
            # This creates natural motion without velocity commands
            alpha = self.smoothness

            next_joints = current_joints + np.clip(alpha * (target - current_joints) , -self.max_ang_inc, self.max_ang_inc)

            # DIRECT POSITION CONTROL - Set joint positions
            self.data.qpos[:self.n_joints] = next_joints

            # IMPORTANT: Set velocities to zero to prevent physics from moving joints
            # This ensures pure kinematic control
            self.data.qvel[:self.n_joints] = 0

            # Step physics (this will apply gravity, collisions, etc. to other bodies)
            mujoco.mj_step(self.model, self.data)

            # After step, ensure our joint positions are maintained
            # (Physics might have altered them slightly)
            self.data.qpos[:self.n_joints] = next_joints
            self.data.qvel[:self.n_joints] = 0

            # Update feedback
            with self.feedback_lock:
                self.current_feedback_joint = self.data.qpos[:self.n_joints].copy()

            # Sync viewer
            if self.viewer:
                self.viewer.sync()

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

    def move_to_joints(self, target, duration=1.0):
        """
        Move to target joints over specified duration

        Args:
            target: Target joint joints
            duration: Time to complete movement (seconds)
        """
        start_js = self.get_feedback()
        end_js = np.array(target[:self.n_joints])

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

        n_steps = int(duration / self.time_interval)

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

        for step in range(n_steps):
            alpha = (step + 1) / n_steps
            # Use easing for smoother motion
            ease_alpha = 1 - (1 - alpha) ** 2  # Quadratic ease-out
            current_target = start_js + ease_alpha * (end_js - start_js)
            self.send_command(current_target)
            time.sleep(self.time_interval)

        # Ensure exact target
        self.send_command(end_js)
        time.sleep(0.1)

    def wait_until_reached(self, tolerance=0.01, timeout=10.0):
        """
        Wait until robot reaches target position

        Args:
            tolerance: Position error tolerance (radians)
            timeout: Maximum wait time (seconds)
        """
        start_time = time.time()

        while time.time() - start_time < timeout:
            current = self.get_feedback()
            target = self.get_target()
            error = np.max(np.abs(target - current))

            if error < tolerance:
                return True

            time.sleep(0.01)

        return False

    def print_state(self):
        """Print current robot state"""
        joints = self.get_feedback()
        target = self.get_target()
        print("Current joints (rad):", [f"{p:.3f}" for p in joints], "...")
        print("Target joints (rad): ", [f"{t:.3f}" for t in target], "...")


# Demo
def demo_position_control():
    """Demonstrate pure position control"""

    urdf_path = "/home/zl/Downloads/urdf_rm75/RM75-B.urdf"

    if not Path(urdf_path).exists():
        print(f"Error: URDF not found at {urdf_path}")
        return

    print("=" * 60)
    print("Pure Position Control Demo")
    print("=" * 60)

    # Create controller
    robot = MuJoCoPositionController(urdf_path, smoothness=0.05, enable_viewer=True)
    robot.start()
    time.sleep(1)

    print("\n[Test 1] Move joint 1 to 45 degrees")
    robot.send_command([0.785, 0, 0, 0, 0, 0, 0])
    robot.wait_until_reached()
    robot.print_state()
    time.sleep(0.5)

    print("\n[Test 2] Move joint 2 to -30 degrees")
    robot.send_command([0, -0.524, 0, 0, 0, 0, 0])
    robot.wait_until_reached()
    robot.print_state()
    time.sleep(0.5)

    print("\n[Test 3] Move multiple joints simultaneously")
    robot.send_command([0.5, -0.4, 0.3, 0.2, 0.1, 0, 0])
    robot.wait_until_reached()
    robot.print_state()
    time.sleep(0.5)

    print("\n[Test 4] Return home")
    robot.send_command([0, 0, 0, 0, 0, 0, 0])
    robot.wait_until_reached()
    robot.print_state()

    print("\n" + "=" * 60)
    print("✓ All tests passed! Robot is stable and controllable.")
    print("=" * 60)
    print("\nInteractive mode - close viewer to exit")

    try:
        while robot.viewer and robot.viewer.is_running():
            time.sleep(0.1)
    except KeyboardInterrupt:
        pass

    robot.stop()




if __name__ == "__main__":
    demo_position_control()
