solver/tests/mates/test_labeling.py

"""Tests for solver.mates.labeling -- mate-level ground truth labels."""

from __future__ import annotations

import numpy as np

from solver.datagen.types import RigidBody
from solver.mates.labeling import MateAssemblyLabels, MateLabel, label_mate_assembly
from solver.mates.patterns import JointPattern
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)]


# ---------------------------------------------------------------------------
# MateLabel
# ---------------------------------------------------------------------------


class TestMateLabel:
    """MateLabel dataclass."""

    def test_defaults(self) -> None:
        ml = MateLabel(mate_id=0)
        assert ml.is_independent is True
        assert ml.is_redundant is False
        assert ml.is_degenerate is False
        assert ml.pattern is None
        assert ml.issue is None

    def test_to_dict(self) -> None:
        ml = MateLabel(
            mate_id=5,
            is_independent=False,
            is_redundant=True,
            pattern=JointPattern.HINGE,
            issue="redundant",
        )
        d = ml.to_dict()
        assert d["mate_id"] == 5
        assert d["is_redundant"] is True
        assert d["pattern"] == "hinge"
        assert d["issue"] == "redundant"

    def test_to_dict_none_pattern(self) -> None:
        ml = MateLabel(mate_id=0)
        d = ml.to_dict()
        assert d["pattern"] is None


# ---------------------------------------------------------------------------
# MateAssemblyLabels
# ---------------------------------------------------------------------------


class TestMateAssemblyLabels:
    """MateAssemblyLabels dataclass."""

    def test_to_dict_structure(self) -> None:
        """to_dict produces expected keys."""
        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 = label_mate_assembly(bodies, mates)
        d = result.to_dict()
        assert "per_mate" in d
        assert "patterns" in d
        assert "assembly" in d
        assert isinstance(d["per_mate"], list)


# ---------------------------------------------------------------------------
# label_mate_assembly
# ---------------------------------------------------------------------------


class TestLabelMateAssembly:
    """Full labeling pipeline."""

    def test_clean_assembly_no_redundancy(self) -> None:
        """Two bodies with lock mate -> clean, no redundancy."""
        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 = label_mate_assembly(bodies, mates)
        assert isinstance(result, MateAssemblyLabels)
        assert len(result.per_mate) == 1
        ml = result.per_mate[0]
        assert ml.mate_id == 0
        assert ml.is_independent is True
        assert ml.is_redundant is False
        assert ml.issue is None

    def test_redundant_assembly(self) -> None:
        """Two lock mates on same body pair -> one is redundant."""
        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),
            ),
            Mate(
                mate_id=1,
                mate_type=MateType.LOCK,
                ref_a=_make_ref(0, GeometryType.FACE, origin=np.array([1.0, 0.0, 0.0])),
                ref_b=_make_ref(1, GeometryType.FACE, origin=np.array([1.0, 0.0, 0.0])),
            ),
        ]
        result = label_mate_assembly(bodies, mates)
        assert len(result.per_mate) == 2
        redundant_count = sum(1 for ml in result.per_mate if ml.is_redundant)
        # At least one should be redundant
        assert redundant_count >= 1
        assert result.assembly.redundant_count > 0

    def test_hinge_pattern_labeling(self) -> None:
        """Hinge mates get pattern membership."""
        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 = label_mate_assembly(bodies, mates)
        assert len(result.per_mate) == 2
        # Both mates should be part of the hinge pattern
        for ml in result.per_mate:
            assert ml.pattern is JointPattern.HINGE
            assert ml.is_independent is True

    def test_grounded_assembly(self) -> None:
        """Grounded assembly labeling 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 = label_mate_assembly(bodies, mates, ground_body=0)
        assert result.assembly.is_rigid

    def test_empty_mates(self) -> None:
        """No mates -> no per_mate labels, underconstrained."""
        bodies = _make_bodies(2)
        result = label_mate_assembly(bodies, [])
        assert len(result.per_mate) == 0
        assert result.assembly.classification == "underconstrained"

    def test_assembly_classification(self) -> None:
        """Assembly classification is present."""
        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 = label_mate_assembly(bodies, mates)
        assert result.assembly.classification in {
            "well-constrained",
            "overconstrained",
            "underconstrained",
            "mixed",
        }