"""Tests for solver.datagen.jacobian -- Jacobian rank verification."""

from __future__ import annotations

import numpy as np
import pytest

from solver.datagen.jacobian import JacobianVerifier
from solver.datagen.types import Joint, JointType, RigidBody

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


def _two_bodies() -> list[RigidBody]:
    return [
        RigidBody(0, position=np.array([0.0, 0.0, 0.0])),
        RigidBody(1, position=np.array([2.0, 0.0, 0.0])),
    ]


def _three_bodies() -> list[RigidBody]:
    return [
        RigidBody(0, position=np.array([0.0, 0.0, 0.0])),
        RigidBody(1, position=np.array([2.0, 0.0, 0.0])),
        RigidBody(2, position=np.array([4.0, 0.0, 0.0])),
    ]


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestJacobianShape:
    """Verify Jacobian matrix dimensions for each joint type."""

    @pytest.mark.parametrize(
        "joint_type,expected_rows",
        [
            (JointType.FIXED, 6),
            (JointType.REVOLUTE, 5),
            (JointType.CYLINDRICAL, 4),
            (JointType.SLIDER, 5),
            (JointType.BALL, 3),
            (JointType.PLANAR, 3),
            (JointType.SCREW, 5),
            (JointType.UNIVERSAL, 4),
            (JointType.PARALLEL, 3),
            (JointType.PERPENDICULAR, 1),
            (JointType.DISTANCE, 1),
        ],
    )
    def test_row_count(self, joint_type: JointType, expected_rows: int) -> None:
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        joint = Joint(
            joint_id=0,
            body_a=0,
            body_b=1,
            joint_type=joint_type,
            anchor_a=np.array([1.0, 0.0, 0.0]),
            anchor_b=np.array([1.0, 0.0, 0.0]),
            axis=np.array([0.0, 0.0, 1.0]),
        )
        n_added = v.add_joint_constraints(joint)
        assert n_added == expected_rows

        j = v.get_jacobian()
        assert j.shape == (expected_rows, 12)  # 2 bodies * 6 cols


class TestNumericalRank:
    """Numerical rank checks for known configurations."""

    def test_fixed_joint_rank_six(self) -> None:
        """Fixed joint between 2 bodies: rank = 6."""
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        j = Joint(
            joint_id=0,
            body_a=0,
            body_b=1,
            joint_type=JointType.FIXED,
            anchor_a=np.array([1.0, 0.0, 0.0]),
            anchor_b=np.array([1.0, 0.0, 0.0]),
        )
        v.add_joint_constraints(j)
        assert v.numerical_rank() == 6

    def test_revolute_joint_rank_five(self) -> None:
        """Revolute joint between 2 bodies: rank = 5."""
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        j = Joint(
            joint_id=0,
            body_a=0,
            body_b=1,
            joint_type=JointType.REVOLUTE,
            anchor_a=np.array([1.0, 0.0, 0.0]),
            anchor_b=np.array([1.0, 0.0, 0.0]),
            axis=np.array([0.0, 0.0, 1.0]),
        )
        v.add_joint_constraints(j)
        assert v.numerical_rank() == 5

    def test_ball_joint_rank_three(self) -> None:
        """Ball joint between 2 bodies: rank = 3."""
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        j = Joint(
            joint_id=0,
            body_a=0,
            body_b=1,
            joint_type=JointType.BALL,
            anchor_a=np.array([1.0, 0.0, 0.0]),
            anchor_b=np.array([1.0, 0.0, 0.0]),
        )
        v.add_joint_constraints(j)
        assert v.numerical_rank() == 3

    def test_empty_jacobian_rank_zero(self) -> None:
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        assert v.numerical_rank() == 0


class TestParallelAxesDegeneracy:
    """Parallel revolute axes create geometric dependencies."""

    def _four_body_loop(self) -> list[RigidBody]:
        return [
            RigidBody(0, position=np.array([0.0, 0.0, 0.0])),
            RigidBody(1, position=np.array([2.0, 0.0, 0.0])),
            RigidBody(2, position=np.array([2.0, 2.0, 0.0])),
            RigidBody(3, position=np.array([0.0, 2.0, 0.0])),
        ]

    def _loop_joints(self, axes: list[np.ndarray]) -> list[Joint]:
        pairs = [(0, 1, [1, 0, 0]), (1, 2, [2, 1, 0]), (2, 3, [1, 2, 0]), (3, 0, [0, 1, 0])]
        return [
            Joint(
                joint_id=i,
                body_a=a,
                body_b=b,
                joint_type=JointType.REVOLUTE,
                anchor_a=np.array(anc, dtype=float),
                anchor_b=np.array(anc, dtype=float),
                axis=axes[i],
            )
            for i, (a, b, anc) in enumerate(pairs)
        ]

    def test_parallel_has_lower_rank(self) -> None:
        """4-body closed loop: all-parallel revolute axes produce lower
        Jacobian rank than mixed axes due to geometric dependency."""
        bodies = self._four_body_loop()
        z_axis = np.array([0.0, 0.0, 1.0])

        # All axes parallel to Z
        v_par = JacobianVerifier(bodies)
        for j in self._loop_joints([z_axis] * 4):
            v_par.add_joint_constraints(j)
        rank_par = v_par.numerical_rank()

        # Mixed axes
        mixed = [
            np.array([0.0, 0.0, 1.0]),
            np.array([0.0, 1.0, 0.0]),
            np.array([0.0, 0.0, 1.0]),
            np.array([1.0, 0.0, 0.0]),
        ]
        v_mix = JacobianVerifier(bodies)
        for j in self._loop_joints(mixed):
            v_mix.add_joint_constraints(j)
        rank_mix = v_mix.numerical_rank()

        assert rank_par < rank_mix


class TestFindDependencies:
    """Dependency detection."""

    def test_fixed_joint_no_dependencies(self) -> None:
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        j = Joint(
            joint_id=0,
            body_a=0,
            body_b=1,
            joint_type=JointType.FIXED,
            anchor_a=np.array([1.0, 0.0, 0.0]),
            anchor_b=np.array([1.0, 0.0, 0.0]),
        )
        v.add_joint_constraints(j)
        assert v.find_dependencies() == []

    def test_duplicate_fixed_has_dependencies(self) -> None:
        """Two fixed joints on same pair: second is fully dependent."""
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        for jid in range(2):
            v.add_joint_constraints(
                Joint(
                    joint_id=jid,
                    body_a=0,
                    body_b=1,
                    joint_type=JointType.FIXED,
                    anchor_a=np.array([1.0, 0.0, 0.0]),
                    anchor_b=np.array([1.0, 0.0, 0.0]),
                )
            )
        deps = v.find_dependencies()
        assert len(deps) == 6  # Second fixed joint entirely redundant

    def test_empty_no_dependencies(self) -> None:
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        assert v.find_dependencies() == []


class TestRowLabels:
    """Row label metadata."""

    def test_labels_match_rows(self) -> None:
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        j = Joint(
            joint_id=7,
            body_a=0,
            body_b=1,
            joint_type=JointType.REVOLUTE,
            anchor_a=np.array([1.0, 0.0, 0.0]),
            anchor_b=np.array([1.0, 0.0, 0.0]),
            axis=np.array([0.0, 0.0, 1.0]),
        )
        v.add_joint_constraints(j)
        assert len(v.row_labels) == 5
        assert all(lab["joint_id"] == 7 for lab in v.row_labels)


class TestPerpendicularPair:
    """Internal _perpendicular_pair utility."""

    @pytest.mark.parametrize(
        "axis",
        [
            np.array([1.0, 0.0, 0.0]),
            np.array([0.0, 1.0, 0.0]),
            np.array([0.0, 0.0, 1.0]),
            np.array([1.0, 1.0, 1.0]) / np.sqrt(3),
        ],
    )
    def test_orthonormal(self, axis: np.ndarray) -> None:
        bodies = _two_bodies()
        v = JacobianVerifier(bodies)
        t1, t2 = v._perpendicular_pair(axis)

        # All unit length
        np.testing.assert_allclose(np.linalg.norm(t1), 1.0, atol=1e-12)
        np.testing.assert_allclose(np.linalg.norm(t2), 1.0, atol=1e-12)

        # Mutually perpendicular
        np.testing.assert_allclose(np.dot(axis, t1), 0.0, atol=1e-12)
        np.testing.assert_allclose(np.dot(axis, t2), 0.0, atol=1e-12)
        np.testing.assert_allclose(np.dot(t1, t2), 0.0, atol=1e-12)
