"""Tests for solver.mates.patterns -- joint pattern recognition."""

from __future__ import annotations

import numpy as np

from solver.datagen.types import JointType
from solver.mates.patterns import JointPattern, PatternMatch, recognize_patterns
from solver.mates.primitives import GeometryRef, GeometryType, Mate, MateType

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


def _make_ref(
    body_id: int,
    geom_type: GeometryType,
    *,
    geometry_id: str = "Geom001",
    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=geometry_id,
        origin=origin,
        direction=direction,
    )


def _make_mate(
    mate_id: int,
    mate_type: MateType,
    body_a: int,
    body_b: int,
    geom_a: GeometryType = GeometryType.FACE,
    geom_b: GeometryType = GeometryType.FACE,
) -> Mate:
    """Factory for Mate with body pair and geometry types."""
    return Mate(
        mate_id=mate_id,
        mate_type=mate_type,
        ref_a=_make_ref(body_a, geom_a),
        ref_b=_make_ref(body_b, geom_b),
    )


# ---------------------------------------------------------------------------
# JointPattern enum
# ---------------------------------------------------------------------------


class TestJointPattern:
    """JointPattern enum."""

    def test_member_count(self) -> None:
        assert len(JointPattern) == 9

    def test_string_values(self) -> None:
        for jp in JointPattern:
            assert isinstance(jp.value, str)

    def test_access_by_name(self) -> None:
        assert JointPattern["HINGE"] is JointPattern.HINGE


# ---------------------------------------------------------------------------
# PatternMatch
# ---------------------------------------------------------------------------


class TestPatternMatch:
    """PatternMatch dataclass."""

    def test_construction(self) -> None:
        mate = _make_mate(0, MateType.LOCK, 0, 1)
        pm = PatternMatch(
            pattern=JointPattern.FIXED,
            mates=[mate],
            body_a=0,
            body_b=1,
            confidence=1.0,
            equivalent_joint_type=JointType.FIXED,
        )
        assert pm.pattern is JointPattern.FIXED
        assert pm.confidence == 1.0
        assert pm.missing_mates == []

    def test_to_dict(self) -> None:
        mate = _make_mate(5, MateType.LOCK, 0, 1)
        pm = PatternMatch(
            pattern=JointPattern.FIXED,
            mates=[mate],
            body_a=0,
            body_b=1,
            confidence=1.0,
            equivalent_joint_type=JointType.FIXED,
        )
        d = pm.to_dict()
        assert d["pattern"] == "fixed"
        assert d["mate_ids"] == [5]
        assert d["equivalent_joint_type"] == "FIXED"


# ---------------------------------------------------------------------------
# recognize_patterns — canonical patterns
# ---------------------------------------------------------------------------


class TestRecognizeCanonical:
    """Full-confidence canonical pattern recognition."""

    def test_empty_input(self) -> None:
        assert recognize_patterns([]) == []

    def test_hinge(self) -> None:
        """Concentric(axis) + Coincident(plane) -> Hinge."""
        mates = [
            _make_mate(0, MateType.CONCENTRIC, 0, 1, GeometryType.AXIS, GeometryType.AXIS),
            _make_mate(1, MateType.COINCIDENT, 0, 1, GeometryType.PLANE, GeometryType.PLANE),
        ]
        results = recognize_patterns(mates)
        top = results[0]
        assert top.pattern is JointPattern.HINGE
        assert top.confidence == 1.0
        assert top.equivalent_joint_type is JointType.REVOLUTE
        assert top.missing_mates == []

    def test_slider(self) -> None:
        """Coincident(plane) + Parallel(axis) -> Slider."""
        mates = [
            _make_mate(0, MateType.COINCIDENT, 0, 1, GeometryType.PLANE, GeometryType.PLANE),
            _make_mate(1, MateType.PARALLEL, 0, 1, GeometryType.AXIS, GeometryType.AXIS),
        ]
        results = recognize_patterns(mates)
        top = results[0]
        assert top.pattern is JointPattern.SLIDER
        assert top.confidence == 1.0
        assert top.equivalent_joint_type is JointType.SLIDER

    def test_cylinder(self) -> None:
        """Concentric(axis) only -> Cylinder."""
        mates = [
            _make_mate(0, MateType.CONCENTRIC, 0, 1, GeometryType.AXIS, GeometryType.AXIS),
        ]
        results = recognize_patterns(mates)
        # Should match cylinder at confidence 1.0
        cylinder = [r for r in results if r.pattern is JointPattern.CYLINDER]
        assert len(cylinder) >= 1
        assert cylinder[0].confidence == 1.0
        assert cylinder[0].equivalent_joint_type is JointType.CYLINDRICAL

    def test_ball(self) -> None:
        """Coincident(point) -> Ball."""
        mates = [
            _make_mate(0, MateType.COINCIDENT, 0, 1, GeometryType.POINT, GeometryType.POINT),
        ]
        results = recognize_patterns(mates)
        top = results[0]
        assert top.pattern is JointPattern.BALL
        assert top.confidence == 1.0
        assert top.equivalent_joint_type is JointType.BALL

    def test_planar_face(self) -> None:
        """Coincident(face) -> Planar."""
        mates = [
            _make_mate(0, MateType.COINCIDENT, 0, 1, GeometryType.FACE, GeometryType.FACE),
        ]
        results = recognize_patterns(mates)
        top = results[0]
        assert top.pattern is JointPattern.PLANAR
        assert top.confidence == 1.0
        assert top.equivalent_joint_type is JointType.PLANAR

    def test_fixed(self) -> None:
        """Lock -> Fixed."""
        mates = [
            _make_mate(0, MateType.LOCK, 0, 1, GeometryType.FACE, GeometryType.FACE),
        ]
        results = recognize_patterns(mates)
        top = results[0]
        assert top.pattern is JointPattern.FIXED
        assert top.confidence == 1.0
        assert top.equivalent_joint_type is JointType.FIXED


# ---------------------------------------------------------------------------
# recognize_patterns — partial matches
# ---------------------------------------------------------------------------


class TestRecognizePartial:
    """Partial pattern matches and hints."""

    def test_concentric_without_plane_hints_hinge(self) -> None:
        """Concentric alone matches hinge at 0.5 confidence with missing hint."""
        mates = [
            _make_mate(0, MateType.CONCENTRIC, 0, 1, GeometryType.AXIS, GeometryType.AXIS),
        ]
        results = recognize_patterns(mates)
        hinge_matches = [r for r in results if r.pattern is JointPattern.HINGE]
        assert len(hinge_matches) >= 1
        hinge = hinge_matches[0]
        assert hinge.confidence == 0.5
        assert len(hinge.missing_mates) > 0

    def test_coincident_plane_without_parallel_hints_slider(self) -> None:
        """Coincident(plane) alone matches slider at 0.5 confidence."""
        mates = [
            _make_mate(0, MateType.COINCIDENT, 0, 1, GeometryType.PLANE, GeometryType.PLANE),
        ]
        results = recognize_patterns(mates)
        slider_matches = [r for r in results if r.pattern is JointPattern.SLIDER]
        assert len(slider_matches) >= 1
        assert slider_matches[0].confidence == 0.5


# ---------------------------------------------------------------------------
# recognize_patterns — ambiguous / multi-body
# ---------------------------------------------------------------------------


class TestRecognizeAmbiguous:
    """Ambiguous patterns and multi-body-pair assemblies."""

    def test_concentric_matches_both_hinge_and_cylinder(self) -> None:
        """A single concentric mate produces both hinge (partial) and cylinder matches."""
        mates = [
            _make_mate(0, MateType.CONCENTRIC, 0, 1, GeometryType.AXIS, GeometryType.AXIS),
        ]
        results = recognize_patterns(mates)
        patterns = {r.pattern for r in results}
        assert JointPattern.HINGE in patterns
        assert JointPattern.CYLINDER in patterns

    def test_multiple_body_pairs(self) -> None:
        """Mates across different body pairs produce separate pattern matches."""
        mates = [
            _make_mate(0, MateType.LOCK, 0, 1),
            _make_mate(1, MateType.COINCIDENT, 2, 3, GeometryType.POINT, GeometryType.POINT),
        ]
        results = recognize_patterns(mates)
        pairs = {(r.body_a, r.body_b) for r in results}
        assert (0, 1) in pairs
        assert (2, 3) in pairs

    def test_results_sorted_by_confidence(self) -> None:
        """All results should be sorted by confidence descending."""
        mates = [
            _make_mate(0, MateType.CONCENTRIC, 0, 1, GeometryType.AXIS, GeometryType.AXIS),
            _make_mate(1, MateType.LOCK, 2, 3),
        ]
        results = recognize_patterns(mates)
        confidences = [r.confidence for r in results]
        assert confidences == sorted(confidences, reverse=True)

    def test_unknown_pattern(self) -> None:
        """A mate type that matches no rule returns UNKNOWN."""
        mates = [
            _make_mate(0, MateType.ANGLE, 0, 1, GeometryType.FACE, GeometryType.FACE),
        ]
        results = recognize_patterns(mates)
        assert any(r.pattern is JointPattern.UNKNOWN for r in results)

    def test_body_pair_normalization(self) -> None:
        """Mates with reversed body order should be grouped together."""
        mates = [
            _make_mate(0, MateType.CONCENTRIC, 1, 0, GeometryType.AXIS, GeometryType.AXIS),
            _make_mate(1, MateType.COINCIDENT, 0, 1, GeometryType.PLANE, GeometryType.PLANE),
        ]
        results = recognize_patterns(mates)
        hinge_matches = [r for r in results if r.pattern is JointPattern.HINGE]
        assert len(hinge_matches) >= 1
        assert hinge_matches[0].confidence == 1.0