#!/usr/bin/env python3
"""
Simple Position Control with Velocity and Acceleration Limits - WITH IMPROVED CONVERGENCE
"""

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


class SimplePositionController:
    """
    Simple position control with velocity and acceleration limits
    """

    def __init__(self, urdf_path,
                 max_velocity=2.0,  # rad/s
                 max_acceleration=5.0,  # rad/s²
                 control_dt=0.01,  # seconds
                 enable_viewer=True):

        # 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.control_dt = control_dt

        # 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")
        print(f"Velocity limit: {max_velocity} rad/s")
        print(f"Acceleration limit: {max_acceleration} rad/s²")

        # Control limits
        self.max_vel = max_velocity
        self.max_acc = max_acceleration

        # Target and current positions
        self.target_positions = self.data.qpos[:self.n_joints].copy()
        self.current_positions = self.data.qpos[:self.n_joints].copy()

        # For motion limiting
        self.current_command = self.data.qpos[:self.n_joints].copy()
        self.previous_command = self.data.qpos[:self.n_joints].copy()

        # For convergence zone - slow down when close to target
        self.slow_zone_radius = 0.2  # rad - start slowing down within this distance
        self.slow_zone_velocity = 0.3  # rad/s - max velocity in slow zone

        # Thread safety
        self.command_lock = threading.Lock()
        self.feedback_lock = threading.Lock()

        # 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}")

        print("Controller ready\n")

    def start(self):
        """Start 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 started")

    def stop(self):
        """Stop simulation"""
        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"""
        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_positions = cmd

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

    def _simulation_loop(self):
        """Main simulation loop"""
        last_time = time.time()

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

            # Calculate distance to target
            distance_to_target = target - self.current_command

            # Adaptive velocity based on distance (slow down when close)
            for i in range(self.n_joints):
                dist = abs(distance_to_target[i])
                if dist < self.slow_zone_radius:
                    # Slow down proportionally to distance
                    max_vel_for_joint = self.slow_zone_velocity * (dist / self.slow_zone_radius)
                    max_vel_for_joint = max(max_vel_for_joint, 0.05)  # Minimum velocity
                else:
                    max_vel_for_joint = self.max_vel

                # Limit the step size
                max_step = max_vel_for_joint * self.control_dt
                if abs(distance_to_target[i]) > max_step:
                    distance_to_target[i] = np.sign(distance_to_target[i]) * max_step

            # Apply acceleration limit
            delta = distance_to_target
            current_vel = (self.current_command - self.previous_command) / self.control_dt
            desired_vel = delta / self.control_dt
            max_vel_change = self.max_acc * self.control_dt

            vel_diff = desired_vel - current_vel
            if np.any(np.abs(vel_diff) > max_vel_change):
                vel_diff = np.clip(vel_diff, -max_vel_change, max_vel_change)
                delta = (current_vel + vel_diff) * self.control_dt

            # Update command
            next_command = self.current_command + delta

            # Store for next iteration
            self.previous_command = self.current_command.copy()
            self.current_command = next_command.copy()

            # Direct position control
            self.data.qpos[:self.n_joints] = next_command
            self.data.qvel[:self.n_joints] = 0

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

            # Maintain position (physics might change it)
            self.data.qpos[:self.n_joints] = next_command

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

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

            # Maintain control rate
            elapsed = time.time() - last_time
            sleep_time = self.control_dt - elapsed
            if sleep_time > 0:
                time.sleep(sleep_time)
            last_time = time.time()

    def wait_for_convergence(self, tolerance=0.01, timeout=2.0):
        """Wait for robot to converge to target"""
        start_time = time.time()

        while time.time() - start_time < timeout:
            current = self.get_feedback()
            with self.command_lock:
                target = self.target_positions

            error = np.max(np.abs(target - current))

            if error < tolerance:
                return True

            time.sleep(0.01)

        return False


def simple_test():
    """Test with improved convergence"""
    urdf_path = "/home/zl/Downloads/urdf_rm75/RM75-B.urdf"

    robot = SimplePositionController(
        urdf_path,
        max_velocity=2.0,  # 2 rad/s
        max_acceleration=5.0,
        control_dt=0.01,
        enable_viewer=True
    )
    robot.start()
    time.sleep(1)

    print("\n" + "=" * 60)
    print("Testing movement with smooth convergence")
    print("=" * 60)

    # Test 1: Move to 0.8 rad and back
    print("\n[Test 1] Moving to +0.8 rad")
    robot.send_command([0.8, 0, 0, 0, 0, 0, 0])
    robot.wait_for_convergence(tolerance=0.01)
    pos = robot.get_feedback()
    print(f"  Converged to: {pos[0]:.4f} rad (error: {abs(pos[0] - 0.8):.4f})")

    print("\n[Test 2] Moving to -0.8 rad")
    robot.send_command([-0.8, 0, 0, 0, 0, 0, 0])
    robot.wait_for_convergence(tolerance=0.01)
    pos = robot.get_feedback()
    print(f"  Converged to: {pos[0]:.4f} rad (error: {abs(pos[0] + 0.8):.4f})")

    print("\n[Test 3] Moving to home (0 rad)")
    robot.send_command([0, 0, 0, 0, 0, 0, 0])
    robot.wait_for_convergence(tolerance=0.01)
    pos = robot.get_feedback()
    print(f"  Converged to: {pos[0]:.4f} rad (error: {abs(pos[0]):.4f})")

    print("\n[Test 4] Multi-joint movement")
    robot.send_command([0.5, -0.3, 0.2, 0.1, 0, 0, 0])
    robot.wait_for_convergence(tolerance=0.01)
    pos = robot.get_feedback()
    print(f"  Final positions: {[f'{p:.3f}' for p in pos[:4]]}...")

    print("\n[Test 5] Return home")
    robot.send_command([0, 0, 0, 0, 0, 0, 0])
    robot.wait_for_convergence(tolerance=0.01)
    pos = robot.get_feedback()
    print(f"  Home positions: {[f'{p:.3f}' for p in pos[:4]]}...")

    print("\n✓ All tests passed! Robot converges quickly to target.")
    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()


def back_and_forth_test():
    """Back and forth test showing smooth convergence"""
    urdf_path = "/home/zl/Downloads/urdf_rm75/RM75-B.urdf"

    robot = SimplePositionController(
        urdf_path,
        max_velocity=2.0,
        max_acceleration=5.0,
        control_dt=0.01,
        enable_viewer=True
    )
    robot.start()
    time.sleep(1)

    print("\n" + "=" * 60)
    print("Back and forth movement with smooth convergence")
    print("=" * 60)

    targets = [[0, -0.524, 0, 0, 0, 0, 0],
               [0.5, -0.4, 0.3, 0.2, 0.1, 0, 0],
               [0.6, -0.5, 0.4, 0.2, 0, 0, 0],
               [0, 0, 0, 0, 0, 0, 0]
               ]

    for target in targets:
        print(f"\nMoving to {target} rad...")
        start_time = time.time()

        robot.send_command(target)

        # Wait for convergence
        robot.wait_for_convergence(tolerance=0.01, timeout=2.0)

        elapsed = time.time() - start_time
        pos = robot.get_feedback()
        error = sum( abs( np.array(pos)  - np.array(target) ) )

        print(f"  Position: {pos[0]:.4f} rad, Error: {error:.4f} rad, Time: {elapsed:.2f}s")

        # Brief pause to observe
        time.sleep(0.3)

    print("\n✓ Test complete! 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__":
    # Run the back and forth test with smooth convergence
    back_and_forth_test()