"""Joint pattern recognition from mate combinations.

Groups mates by body pair and matches them against canonical joint
patterns (hinge, slider, ball, etc.). Each pattern is a known
combination of mate types that together constrain motion equivalently
to a single mechanical joint.
"""

from __future__ import annotations

import enum
from collections import defaultdict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from solver.datagen.types import JointType
from solver.mates.primitives import GeometryType, Mate, MateType

if TYPE_CHECKING:
    from typing import Any

__all__ = [
    "JointPattern",
    "PatternMatch",
    "recognize_patterns",
]


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------


class JointPattern(enum.Enum):
    """Canonical joint patterns formed by mate combinations."""

    HINGE = "hinge"
    SLIDER = "slider"
    CYLINDER = "cylinder"
    BALL = "ball"
    PLANAR = "planar"
    FIXED = "fixed"
    GEAR = "gear"
    RACK_PINION = "rack_pinion"
    UNKNOWN = "unknown"


# ---------------------------------------------------------------------------
# Pattern match result
# ---------------------------------------------------------------------------


@dataclass
class PatternMatch:
    """Result of matching a group of mates to a joint pattern.

    Attributes:
        pattern: The identified joint pattern.
        mates: The mates that form this pattern.
        body_a: First body in the pair.
        body_b: Second body in the pair.
        confidence: How well the mates match the canonical pattern (0-1).
        equivalent_joint_type: The JointType this pattern maps to.
        missing_mates: Descriptions of mates absent for a full match.
    """

    pattern: JointPattern
    mates: list[Mate]
    body_a: int
    body_b: int
    confidence: float
    equivalent_joint_type: JointType
    missing_mates: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable dict."""
        return {
            "pattern": self.pattern.value,
            "body_a": self.body_a,
            "body_b": self.body_b,
            "confidence": self.confidence,
            "equivalent_joint_type": self.equivalent_joint_type.name,
            "mate_ids": [m.mate_id for m in self.mates],
            "missing_mates": self.missing_mates,
        }


# ---------------------------------------------------------------------------
# Pattern rules (data-driven)
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class _MateRequirement:
    """A single mate requirement within a pattern rule."""

    mate_type: MateType
    geometry_a: GeometryType | None = None
    geometry_b: GeometryType | None = None


@dataclass(frozen=True)
class _PatternRule:
    """Defines a canonical pattern as a set of required mates."""

    pattern: JointPattern
    joint_type: JointType
    required: tuple[_MateRequirement, ...]
    description: str = ""


_PATTERN_RULES: list[_PatternRule] = [
    _PatternRule(
        pattern=JointPattern.HINGE,
        joint_type=JointType.REVOLUTE,
        required=(
            _MateRequirement(MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS),
            _MateRequirement(MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE),
        ),
        description="Concentric axes + coincident plane",
    ),
    _PatternRule(
        pattern=JointPattern.SLIDER,
        joint_type=JointType.SLIDER,
        required=(
            _MateRequirement(MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE),
            _MateRequirement(MateType.PARALLEL, GeometryType.AXIS, GeometryType.AXIS),
        ),
        description="Coincident plane + parallel axis",
    ),
    _PatternRule(
        pattern=JointPattern.CYLINDER,
        joint_type=JointType.CYLINDRICAL,
        required=(_MateRequirement(MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS),),
        description="Concentric axes only",
    ),
    _PatternRule(
        pattern=JointPattern.BALL,
        joint_type=JointType.BALL,
        required=(_MateRequirement(MateType.COINCIDENT, GeometryType.POINT, GeometryType.POINT),),
        description="Coincident points",
    ),
    _PatternRule(
        pattern=JointPattern.PLANAR,
        joint_type=JointType.PLANAR,
        required=(_MateRequirement(MateType.COINCIDENT, GeometryType.FACE, GeometryType.FACE),),
        description="Coincident faces",
    ),
    _PatternRule(
        pattern=JointPattern.PLANAR,
        joint_type=JointType.PLANAR,
        required=(_MateRequirement(MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE),),
        description="Coincident planes (alternate planar)",
    ),
    _PatternRule(
        pattern=JointPattern.FIXED,
        joint_type=JointType.FIXED,
        required=(_MateRequirement(MateType.LOCK),),
        description="Lock mate",
    ),
]


# ---------------------------------------------------------------------------
# Matching logic
# ---------------------------------------------------------------------------


def _mate_matches_requirement(mate: Mate, req: _MateRequirement) -> bool:
    """Check if a mate satisfies a requirement."""
    if mate.mate_type is not req.mate_type:
        return False
    if req.geometry_a is not None and mate.ref_a.geometry_type is not req.geometry_a:
        return False
    return not (req.geometry_b is not None and mate.ref_b.geometry_type is not req.geometry_b)


def _try_match_rule(
    rule: _PatternRule,
    mates: list[Mate],
) -> tuple[float, list[Mate], list[str]]:
    """Try to match a rule against a group of mates.

    Returns:
        (confidence, matched_mates, missing_descriptions)
    """
    matched: list[Mate] = []
    missing: list[str] = []

    for req in rule.required:
        found = False
        for mate in mates:
            if mate in matched:
                continue
            if _mate_matches_requirement(mate, req):
                matched.append(mate)
                found = True
                break
        if not found:
            geom_desc = ""
            if req.geometry_a is not None:
                geom_b = req.geometry_b.value if req.geometry_b else "*"
                geom_desc = f" ({req.geometry_a.value}-{geom_b})"
            missing.append(f"{req.mate_type.name}{geom_desc}")

    total_required = len(rule.required)
    if total_required == 0:
        return 0.0, [], []

    matched_count = len(matched)
    confidence = matched_count / total_required

    return confidence, matched, missing


def _normalize_body_pair(body_a: int, body_b: int) -> tuple[int, int]:
    """Normalize a body pair so the smaller ID comes first."""
    return (min(body_a, body_b), max(body_a, body_b))


def recognize_patterns(mates: list[Mate]) -> list[PatternMatch]:
    """Identify joint patterns from a list of mates.

    Groups mates by body pair, then checks each group against
    canonical pattern rules. Returns matches sorted by confidence
    descending.

    Args:
        mates: List of mate constraints to analyze.

    Returns:
        List of PatternMatch results, highest confidence first.
    """
    if not mates:
        return []

    # Group mates by normalized body pair
    groups: dict[tuple[int, int], list[Mate]] = defaultdict(list)
    for mate in mates:
        pair = _normalize_body_pair(mate.ref_a.body_id, mate.ref_b.body_id)
        groups[pair].append(mate)

    results: list[PatternMatch] = []

    for (body_a, body_b), group_mates in groups.items():
        group_matches: list[PatternMatch] = []

        for rule in _PATTERN_RULES:
            confidence, matched, missing = _try_match_rule(rule, group_mates)

            if confidence > 0:
                group_matches.append(
                    PatternMatch(
                        pattern=rule.pattern,
                        mates=matched if matched else group_mates,
                        body_a=body_a,
                        body_b=body_b,
                        confidence=confidence,
                        equivalent_joint_type=rule.joint_type,
                        missing_mates=missing,
                    )
                )

        if group_matches:
            # Sort by confidence descending, prefer more-specific patterns
            group_matches.sort(key=lambda m: (-m.confidence, -len(m.mates)))
            results.extend(group_matches)
        else:
            # No pattern matched at all
            results.append(
                PatternMatch(
                    pattern=JointPattern.UNKNOWN,
                    mates=group_mates,
                    body_a=body_a,
                    body_b=body_b,
                    confidence=0.0,
                    equivalent_joint_type=JointType.DISTANCE,
                    missing_mates=[],
                )
            )

    # Global sort by confidence descending
    results.sort(key=lambda m: -m.confidence)
    return results
