#!/usr/bin/env python3
"""Tests for joint limit recovery velocity bounds."""

import os
import tempfile

import numpy as np
import pytest

import embodik as eik


def create_test_urdf():
    """Create a simple 1-DoF URDF with joint limits."""
    urdf_content = """<?xml version="1.0"?>
<robot name="test_robot">
  <link name="base_link">
    <inertial>
      <mass value="1.0"/>
      <origin xyz="0 0 0"/>
      <inertia ixx="1" ixy="0" ixz="0" iyy="1" iyz="0" izz="1"/>
    </inertial>
  </link>

  <link name="link1">
    <inertial>
      <mass value="1.0"/>
      <origin xyz="0 0 0.5"/>
      <inertia ixx="1" ixy="0" ixz="0" iyy="1" iyz="1" izz="1"/>
    </inertial>
  </link>

  <joint name="joint1" type="revolute">
    <parent link="base_link"/>
    <child link="link1"/>
    <origin xyz="0 0 1" rpy="0 0 0"/>
    <axis xyz="0 0 1"/>
    <limit lower="0.0" upper="1.0" velocity="1.0" effort="10.0"/>
  </joint>
</robot>"""
    fd, path = tempfile.mkstemp(suffix=".urdf")
    with os.fdopen(fd, "w") as f:
        f.write(urdf_content)
    return path


def create_continuous_plus_revolute_urdf():
    """Create a 2-joint chain with continuous then bounded revolute joint."""
    urdf_content = """<?xml version="1.0"?>
<robot name="mixed_joint_robot">
  <link name="base_link"/>
  <link name="link1"/>
  <link name="link2"/>

  <joint name="joint_cont" type="continuous">
    <parent link="base_link"/>
    <child link="link1"/>
    <origin xyz="0 0 0" rpy="0 0 0"/>
    <axis xyz="0 0 1"/>
    <limit velocity="2.0" effort="10.0"/>
  </joint>

  <joint name="joint_limited" type="revolute">
    <parent link="link1"/>
    <child link="link2"/>
    <origin xyz="0 0 1" rpy="0 0 0"/>
    <axis xyz="0 1 0"/>
    <limit lower="0.0" upper="1.0" velocity="2.0" effort="10.0"/>
  </joint>
</robot>"""
    fd, path = tempfile.mkstemp(suffix=".urdf")
    with os.fdopen(fd, "w") as f:
        f.write(urdf_content)
    return path


@pytest.fixture
def solver():
    urdf_path = create_test_urdf()
    model = eik.RobotModel(urdf_path, floating_base=False)
    solver = eik.KinematicsSolver(model)
    yield solver
    os.remove(urdf_path)


@pytest.fixture
def constraint_params():
    return {"velocity_limit": 1.0, "acceleration_limit": 10.0, "dt": 0.01}


class TestVelocityBoxConstraintFormulation:
    """Direct tests of the velocity box constraint math."""

    def test_current_inside_limits(self, solver, constraint_params):
        """Inside limits: symmetric velocity bounds."""
        lower, upper = solver.calculate_velocity_box_constraint(
            0.5,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert np.isclose(lower, -1.0)
        assert np.isclose(upper, 1.0)

    def test_current_at_lower_limit(self, solver, constraint_params):
        """At lower limit: softening does NOT inject velocity toward the limit."""
        lower, upper = solver.calculate_velocity_box_constraint(
            -1e-4,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert lower <= 0.0
        assert lower >= -0.15
        assert np.isclose(upper, 1.0)

    def test_new_outside_lower_forces_recovery(self, solver, constraint_params):
        """Outside lower limit: v >= v_recovery required."""
        solver.set_limit_recovery_gain(0.5)
        lower, upper = solver.calculate_velocity_box_constraint(
            -0.0501,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert np.isclose(lower, 1.0)
        assert upper >= lower

    def test_new_outside_upper_forces_recovery(self, solver, constraint_params):
        """Outside upper limit: v <= -v_recovery required."""
        solver.set_limit_recovery_gain(0.5)
        lower, upper = solver.calculate_velocity_box_constraint(
            0.5,
            -0.0501,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert np.isclose(upper, -1.0)
        assert lower <= upper

    def test_recovery_gain_affects_speed(self, solver, constraint_params):
        """Different recovery gains produce different recovery velocities."""
        solver.set_limit_recovery_gain(0.1)
        lower_gentle, _ = solver.calculate_velocity_box_constraint(
            -0.0201,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        solver.set_limit_recovery_gain(1.0)
        lower_aggressive, _ = solver.calculate_velocity_box_constraint(
            -0.0201,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert lower_gentle < lower_aggressive

    def test_small_violation_small_recovery(self, solver, constraint_params):
        """Small violations need small recovery velocities."""
        solver.set_limit_recovery_gain(0.5)
        lower, _ = solver.calculate_velocity_box_constraint(
            -0.0011,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert np.isclose(lower, 0.055)

    def test_release_margin_reduces_tiny_violation_recovery(self, solver, constraint_params):
        """Release margin should reduce forced recovery near boundary."""
        solver.set_limit_recovery_gain(0.5)
        base_lower, _ = solver.calculate_velocity_box_constraint(
            -0.0011,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        solver.set_limit_exit_release_margin(0.0010)
        reduced_lower, _ = solver.calculate_velocity_box_constraint(
            -0.0011,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert reduced_lower < base_lower
        assert np.isclose(reduced_lower, 0.005)

    def test_recovery_hysteresis_widens_no_force_zone(self, solver, constraint_params):
        """Larger enter epsilon should avoid forced recovery for tiny violations."""
        solver.set_limit_recovery_gain(0.5)
        lower_default, _ = solver.calculate_velocity_box_constraint(
            -0.0002,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        solver.set_limit_recovery_hysteresis(3e-4, 5e-4)
        lower_hysteresis, _ = solver.calculate_velocity_box_constraint(
            -0.0002,
            0.5,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert lower_default > 0.0
        assert lower_hysteresis <= 0.0

    def test_multi_step_convergence(self, solver, constraint_params):
        """Verify recovery to valid region over multiple iterations."""
        solver.set_limit_recovery_gain(0.5)
        q = -0.1
        q_min = 0.0
        q_max = 1.0
        margin_limit = 1e-4
        for _ in range(40):
            lower_margin = q - q_min - margin_limit
            upper_margin = q_max - q - margin_limit
            lower, _ = solver.calculate_velocity_box_constraint(
                lower_margin,
                upper_margin,
                constraint_params["velocity_limit"],
                constraint_params["acceleration_limit"],
                constraint_params["dt"],
            )
            v = lower if lower > 0 else 0.0
            q += v * constraint_params["dt"]
        assert q >= q_min - 1e-4

    def test_no_sign_flip_chatter_during_limit_recovery(self, solver, constraint_params):
        """Recovery from a lower-limit violation should be monotonic."""
        solver.set_limit_recovery_gain(0.5)
        solver.set_limit_recovery_hysteresis(1e-4, 3e-4)
        solver.set_limit_exit_release_margin(5e-4)

        q = -0.01
        q_min = 0.0
        q_max = 1.0
        margin_limit = 1e-4
        previous_velocity = None
        sign_flip_count = 0

        for _ in range(15):
            lower_margin = q - q_min - margin_limit
            upper_margin = q_max - q - margin_limit
            lower, upper = solver.calculate_velocity_box_constraint(
                lower_margin,
                upper_margin,
                constraint_params["velocity_limit"],
                constraint_params["acceleration_limit"],
                constraint_params["dt"],
            )
            v = np.clip(1.0, lower, upper)
            if previous_velocity is not None and previous_velocity > 0.0 and v < 0.0:
                sign_flip_count += 1
            previous_velocity = v
            q += v * constraint_params["dt"]

        assert sign_flip_count == 0
        assert q >= q_min - 1e-4

    def test_both_limits_violated_infeasible(self, solver, constraint_params):
        """If both margins are negative, constraint is infeasible."""
        lower, upper = solver.calculate_velocity_box_constraint(
            -0.1,
            -0.2,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert np.isclose(lower, -1.0)
        assert np.isclose(upper, 1.0)

    def test_exactly_at_boundary(self, solver, constraint_params):
        """At exact boundary, softening does NOT inject velocity toward limit."""
        lower, upper = solver.calculate_velocity_box_constraint(
            -1e-4,
            0.3,
            constraint_params["velocity_limit"],
            constraint_params["acceleration_limit"],
            constraint_params["dt"],
        )
        assert lower <= 0.0
        assert lower >= -0.15
        assert upper > 0.0


def test_mixed_joint_index_mapping_preserves_revolute_limit():
    """Regression: position limits must hold with continuous joints in model."""
    urdf_path = create_continuous_plus_revolute_urdf()
    try:
        robot = eik.RobotModel(urdf_path, floating_base=False)
        solver = eik.KinematicsSolver(robot)
        solver.dt = 0.1

        # Drive bounded joint toward +upper limit aggressively.
        task = solver.add_joint_task("bounded_joint_task", "joint_limited", 2.0)
        task.weight = 1.0

        q = robot.get_current_configuration().copy()
        # joint_cont occupies q[0:2] as [cos, sin]; bounded revolute is q[2].
        q[2] = 0.99
        robot.update_configuration(q)

        result = solver.solve_velocity(q, apply_limits=True)
        assert result.status == eik.SolverStatus.SUCCESS

        q_next = robot.integrate(q, np.asarray(result.joint_velocities), solver.dt)
        q_lower, q_upper = robot.get_joint_limits()
        q_next_clamped = np.clip(q_next, q_lower, q_upper)
        assert q_next_clamped[2] <= 1.0 + 1e-8
    finally:
        os.remove(urdf_path)
