"""Tests for solver.mates.conversion -- mate-to-joint conversion."""

from __future__ import annotations

import numpy as np

from solver.datagen.types import JointType, RigidBody
from solver.mates.conversion import (
    MateAnalysisResult,
    analyze_mate_assembly,
    convert_mates_to_joints,
)
from solver.mates.primitives import GeometryRef, GeometryType, Mate, MateType

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_ref(
    body_id: int,
    geom_type: GeometryType,
    *,
    origin: np.ndarray | None = None,
    direction: np.ndarray | None = None,
) -> GeometryRef:
    """Factory for GeometryRef with sensible defaults."""
    if origin is None:
        origin = np.zeros(3)
    if direction is None and geom_type in {
        GeometryType.FACE,
        GeometryType.AXIS,
        GeometryType.PLANE,
    }:
        direction = np.array([0.0, 0.0, 1.0])
    return GeometryRef(
        body_id=body_id,
        geometry_type=geom_type,
        geometry_id="Geom001",
        origin=origin,
        direction=direction,
    )


def _make_bodies(n: int) -> list[RigidBody]:
    """Create n bodies at distinct positions."""
    return [RigidBody(body_id=i, position=np.array([float(i), 0.0, 0.0])) for i in range(n)]


# ---------------------------------------------------------------------------
# convert_mates_to_joints
# ---------------------------------------------------------------------------


class TestConvertMatesToJoints:
    """convert_mates_to_joints function."""

    def test_empty_input(self) -> None:
        joints, m2j, j2m = convert_mates_to_joints([])
        assert joints == []
        assert m2j == {}
        assert j2m == {}

    def test_hinge_pattern(self) -> None:
        """Concentric + Coincident(plane) -> single REVOLUTE joint."""
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.CONCENTRIC,
                ref_a=_make_ref(0, GeometryType.AXIS),
                ref_b=_make_ref(1, GeometryType.AXIS),
            ),
            Mate(
                mate_id=1,
                mate_type=MateType.COINCIDENT,
                ref_a=_make_ref(0, GeometryType.PLANE),
                ref_b=_make_ref(1, GeometryType.PLANE),
            ),
        ]
        joints, m2j, j2m = convert_mates_to_joints(mates)
        assert len(joints) == 1
        assert joints[0].joint_type is JointType.REVOLUTE
        assert joints[0].body_a == 0
        assert joints[0].body_b == 1
        # Both mates map to the single joint
        assert 0 in m2j
        assert 1 in m2j
        assert j2m[joints[0].joint_id] == [0, 1]

    def test_lock_pattern(self) -> None:
        """Lock -> FIXED joint."""
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.LOCK,
                ref_a=_make_ref(0, GeometryType.FACE),
                ref_b=_make_ref(1, GeometryType.FACE),
            ),
        ]
        joints, _m2j, _j2m = convert_mates_to_joints(mates)
        assert len(joints) == 1
        assert joints[0].joint_type is JointType.FIXED

    def test_unmatched_mate_fallback(self) -> None:
        """A single ANGLE mate with no pattern -> individual joint."""
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.ANGLE,
                ref_a=_make_ref(0, GeometryType.FACE),
                ref_b=_make_ref(1, GeometryType.FACE),
            ),
        ]
        joints, _m2j, _j2m = convert_mates_to_joints(mates)
        assert len(joints) == 1
        assert joints[0].joint_type is JointType.PERPENDICULAR

    def test_mapping_consistency(self) -> None:
        """mate_to_joint and joint_to_mates are consistent."""
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.CONCENTRIC,
                ref_a=_make_ref(0, GeometryType.AXIS),
                ref_b=_make_ref(1, GeometryType.AXIS),
            ),
            Mate(
                mate_id=1,
                mate_type=MateType.COINCIDENT,
                ref_a=_make_ref(0, GeometryType.PLANE),
                ref_b=_make_ref(1, GeometryType.PLANE),
            ),
            Mate(
                mate_id=2,
                mate_type=MateType.DISTANCE,
                ref_a=_make_ref(2, GeometryType.POINT),
                ref_b=_make_ref(3, GeometryType.POINT),
            ),
        ]
        joints, m2j, j2m = convert_mates_to_joints(mates)
        # Every mate should be in m2j
        for mate in mates:
            assert mate.mate_id in m2j
        # Every joint should be in j2m
        for joint in joints:
            assert joint.joint_id in j2m

    def test_joint_axis_from_geometry(self) -> None:
        """Joint axis should come from mate geometry direction."""
        axis_dir = np.array([1.0, 0.0, 0.0])
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.CONCENTRIC,
                ref_a=_make_ref(0, GeometryType.AXIS, direction=axis_dir),
                ref_b=_make_ref(1, GeometryType.AXIS, direction=axis_dir),
            ),
            Mate(
                mate_id=1,
                mate_type=MateType.COINCIDENT,
                ref_a=_make_ref(0, GeometryType.PLANE),
                ref_b=_make_ref(1, GeometryType.PLANE),
            ),
        ]
        joints, _, _ = convert_mates_to_joints(mates)
        np.testing.assert_array_almost_equal(joints[0].axis, axis_dir)


# ---------------------------------------------------------------------------
# MateAnalysisResult
# ---------------------------------------------------------------------------


class TestMateAnalysisResult:
    """MateAnalysisResult dataclass."""

    def test_to_dict(self) -> None:
        result = MateAnalysisResult(
            patterns=[],
            joints=[],
        )
        d = result.to_dict()
        assert d["patterns"] == []
        assert d["joints"] == []
        assert d["labels"] is None


# ---------------------------------------------------------------------------
# analyze_mate_assembly
# ---------------------------------------------------------------------------


class TestAnalyzeMateAssembly:
    """Full pipeline: mates -> joints -> analysis."""

    def test_two_bodies_hinge(self) -> None:
        """Two bodies connected by hinge mates -> underconstrained (1 DOF)."""
        bodies = _make_bodies(2)
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.CONCENTRIC,
                ref_a=_make_ref(0, GeometryType.AXIS),
                ref_b=_make_ref(1, GeometryType.AXIS),
            ),
            Mate(
                mate_id=1,
                mate_type=MateType.COINCIDENT,
                ref_a=_make_ref(0, GeometryType.PLANE),
                ref_b=_make_ref(1, GeometryType.PLANE),
            ),
        ]
        result = analyze_mate_assembly(bodies, mates)
        assert result.analysis is not None
        assert result.labels is not None
        # A revolute joint removes 5 DOF, leaving 1 internal DOF
        assert result.analysis.combinatorial_internal_dof == 1
        assert len(result.joints) == 1
        assert result.joints[0].joint_type is JointType.REVOLUTE

    def test_two_bodies_fixed(self) -> None:
        """Two bodies with lock mate -> well-constrained."""
        bodies = _make_bodies(2)
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.LOCK,
                ref_a=_make_ref(0, GeometryType.FACE),
                ref_b=_make_ref(1, GeometryType.FACE),
            ),
        ]
        result = analyze_mate_assembly(bodies, mates)
        assert result.analysis is not None
        assert result.analysis.combinatorial_internal_dof == 0
        assert result.analysis.is_rigid

    def test_grounded_assembly(self) -> None:
        """Grounded assembly analysis works."""
        bodies = _make_bodies(2)
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.LOCK,
                ref_a=_make_ref(0, GeometryType.FACE),
                ref_b=_make_ref(1, GeometryType.FACE),
            ),
        ]
        result = analyze_mate_assembly(bodies, mates, ground_body=0)
        assert result.analysis is not None
        assert result.analysis.is_rigid

    def test_no_mates(self) -> None:
        """Assembly with no mates should be fully underconstrained."""
        bodies = _make_bodies(2)
        result = analyze_mate_assembly(bodies, [])
        assert result.analysis is not None
        assert result.analysis.combinatorial_internal_dof == 6
        assert len(result.joints) == 0

    def test_single_body(self) -> None:
        """Single body, no mates."""
        bodies = _make_bodies(1)
        result = analyze_mate_assembly(bodies, [])
        assert result.analysis is not None
        assert len(result.joints) == 0

    def test_result_traceability(self) -> None:
        """mate_to_joint and joint_to_mates populated in result."""
        bodies = _make_bodies(2)
        mates = [
            Mate(
                mate_id=0,
                mate_type=MateType.CONCENTRIC,
                ref_a=_make_ref(0, GeometryType.AXIS),
                ref_b=_make_ref(1, GeometryType.AXIS),
            ),
            Mate(
                mate_id=1,
                mate_type=MateType.COINCIDENT,
                ref_a=_make_ref(0, GeometryType.PLANE),
                ref_b=_make_ref(1, GeometryType.PLANE),
            ),
        ]
        result = analyze_mate_assembly(bodies, mates)
        assert 0 in result.mate_to_joint
        assert 1 in result.mate_to_joint
        assert len(result.joint_to_mates) > 0
