"""Regression tests for the bounded SEW retargeting differential."""

from pathlib import Path
import sys
import unittest

import numpy as np
import pinocchio as pin


CODE_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(CODE_ROOT))

from core.sew_mapper2 import (  # noqa: E402
    BallJointConfig,
    SEWMapper,
    _wrap_angle_delta,
)


def build_real_mapper() -> SEWMapper:
    config_dir = CODE_ROOT / "config"
    master_model = pin.buildModelFromUrdf(str(config_dir / "master_7dof.urdf"))
    slave_model = pin.buildModelFromUrdf(str(config_dir / "real_slave_7dof.urdf"))

    master_joints = (
        "master_shoulder_pitch_joint",
        "master_shoulder_yaw_joint",
        "master_shoulder_roll_joint",
        "master_elbow_flex_joint",
        "master_wrist_roll_joint",
        "master_wrist_yaw_joint",
        "master_wrist_pitch_joint",
    )
    slave_joints = (
        "R_SHOULDER_P",
        "R_SHOULDER_R",
        "R_SHOULDER_Y",
        "R_ELBOW_R",
        "R_WRIST_P",
        "R_WRIST_Y",
        "R_WRIST_R",
    )
    shoulder = BallJointConfig(
        axis_order="yxy",
        joint_names=("R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y"),
        signs=(-1.0, 1.0, -1.0),
    )
    wrist = BallJointConfig(
        axis_order="yzx",
        joint_names=("R_WRIST_P", "R_WRIST_Y", "R_WRIST_R"),
        signs=(-1.0, 1.0, 1.0),
    )
    return SEWMapper(
        master_model=master_model,
        slave_model=slave_model,
        m_shoulder="master_shoulder",
        m_elbow="master_forearm",
        m_wrist="master_wrist",
        m_ee="master_ee",
        s_shoulder="R_SHOULDER_R_S",
        s_elbow="R_ELBOW_R_S",
        s_wrist="R_WRIST_R_S",
        s_ee="R_WRIST_R_S",
        master_joint_names=master_joints,
        slave_joint_names=slave_joints,
        slave_shoulder_cfg=shoulder,
        slave_wrist_cfg=wrist,
    )


class TestRetargetingDifferential(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.mapper = build_real_mapper()
        # An interior, exactly recoverable pose on the real slave arm.
        cls.q_master = np.array(
            [0.534, 0.314, -0.1, 2.14, 0.38, 0.38, -0.72],
            dtype=float,
        )

    def test_virtual_work_identity(self) -> None:
        q_slave, A, debug = self.mapper.retarget_with_differential(self.q_master)

        self.assertTrue(debug["success"], debug["events"])
        self.assertTrue(
            debug["differential_valid"],
            debug["differential"]["events"],
        )
        self.assertFalse(debug["clipped"])
        self.assertFalse(debug["near_limit"])
        q_slave7 = self.mapper._slave_q7(q_slave)
        self.assertTrue(np.all(q_slave7 >= self.mapper.s_lower7))
        self.assertTrue(np.all(q_slave7 <= self.mapper.s_upper7))

        qdot_master = np.array(
            [0.21, -0.34, 0.17, 0.09, -0.11, 0.28, -0.07]
        )
        tau_slave = np.array(
            [-1.2, 0.4, 2.3, -0.7, 0.5, -1.1, 0.8]
        )
        qdot_slave = A @ qdot_master
        tau_master = A.T @ tau_slave
        self.assertAlmostEqual(
            float(tau_master @ qdot_master),
            float(tau_slave @ qdot_slave),
            places=12,
        )

    def test_matches_independent_directional_finite_difference(self) -> None:
        q_slave, A, debug = self.mapper.retarget_with_differential(
            self.q_master, fd_step=1e-4
        )
        self.assertTrue(debug["differential_valid"])

        direction = np.array([0.3, -0.2, 0.5, -0.1, 0.4, -0.25, 0.35])
        direction /= np.linalg.norm(direction)
        independent_step = 2.5e-5
        q_plus, plus_debug = self.mapper.retarget(
            self.q_master + independent_step * direction,
            q_s_init=q_slave,
        )
        q_minus, minus_debug = self.mapper.retarget(
            self.q_master - independent_step * direction,
            q_s_init=q_slave,
        )
        self.assertTrue(plus_debug["success"], plus_debug["events"])
        self.assertTrue(minus_debug["success"], minus_debug["events"])

        directional_fd = _wrap_angle_delta(
            self.mapper._slave_q7(q_plus) - self.mapper._slave_q7(q_minus)
        ) / (2.0 * independent_step)
        np.testing.assert_allclose(
            directional_fd,
            A @ direction,
            rtol=2e-4,
            atol=2e-5,
        )

    def test_degenerate_geometry_invalidates_differential(self) -> None:
        # Zero elbow flex makes the master shoulder-elbow-wrist plane undefined.
        q_degenerate = np.array([-np.pi / 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        A, debug = self.mapper.compute_differential(q_degenerate)

        self.assertFalse(debug["valid"])
        self.assertTrue(np.isnan(A).all())
        self.assertIn("base_invalid", debug["events"])
        self.assertIn(
            "master_arm_plane_degenerate",
            debug["base"]["events"],
        )

    def test_active_limit_invalidates_differential(self) -> None:
        old_margin = self.mapper.joint_limit_margin
        try:
            # Enlarge only the diagnostic margin to exercise the active-set
            # safety gate without changing the bounded pose solution.
            self.mapper.joint_limit_margin = 10.0
            A, debug = self.mapper.compute_differential(self.q_master)
        finally:
            self.mapper.joint_limit_margin = old_margin

        self.assertTrue(debug["base"]["success"])
        self.assertTrue(debug["base"]["near_limit"])
        self.assertFalse(debug["valid"])
        self.assertTrue(np.isnan(A).all())
        self.assertIn("base_nonsmooth", debug["events"])

    def test_angle_difference_wrap(self) -> None:
        raw = np.array([2.0 * np.pi - 0.1, -2.0 * np.pi + 0.2, np.pi])
        np.testing.assert_allclose(
            _wrap_angle_delta(raw),
            np.array([-0.1, 0.2, -np.pi]),
            atol=1e-14,
        )


if __name__ == "__main__":
    unittest.main()