"""Mate type definitions and geometry references for assembly constraints.

Mates are the user-facing constraint primitives in CAD (e.g. SolidWorks-style
Coincident, Concentric, Parallel). Each mate references geometry on two bodies
and removes a context-dependent number of degrees of freedom.
"""

from __future__ import annotations

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

import numpy as np

if TYPE_CHECKING:
    from typing import Any

__all__ = [
    "GeometryRef",
    "GeometryType",
    "Mate",
    "MateType",
    "dof_removed",
]


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


class MateType(enum.Enum):
    """CAD mate types with default DOF-removal counts.

    Values are ``(ordinal, default_dof)`` tuples so that mate types
    sharing the same DOF count remain distinct enum members. Use the
    :attr:`default_dof` property to get the scalar constraint count.

    The actual DOF removed can be context-dependent (e.g. COINCIDENT
    removes 3 DOF for face-face but only 1 for face-point). Use
    :func:`dof_removed` for the context-aware count.
    """

    COINCIDENT = (0, 3)
    CONCENTRIC = (1, 2)
    PARALLEL = (2, 2)
    PERPENDICULAR = (3, 1)
    TANGENT = (4, 1)
    DISTANCE = (5, 1)
    ANGLE = (6, 1)
    LOCK = (7, 6)

    @property
    def default_dof(self) -> int:
        """Default number of DOF removed by this mate type."""
        return self.value[1]


class GeometryType(enum.Enum):
    """Types of geometric references used by mates."""

    FACE = "face"
    EDGE = "edge"
    POINT = "point"
    AXIS = "axis"
    PLANE = "plane"


# Geometry types that require a direction vector.
_DIRECTIONAL_TYPES = frozenset(
    {
        GeometryType.FACE,
        GeometryType.AXIS,
        GeometryType.PLANE,
    }
)


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------


@dataclass
class GeometryRef:
    """A reference to a specific geometric entity on a body.

    Attributes:
        body_id: Index of the body this geometry belongs to.
        geometry_type: What kind of geometry (face, edge, etc.).
        geometry_id: CAD identifier string (e.g. ``"Face001"``).
        origin: 3D position of the geometry reference point.
        direction: Unit direction vector. Required for FACE, AXIS, PLANE;
            ``None`` for POINT.
    """

    body_id: int
    geometry_type: GeometryType
    geometry_id: str
    origin: np.ndarray = field(default_factory=lambda: np.zeros(3))
    direction: np.ndarray | None = None

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable dict."""
        return {
            "body_id": self.body_id,
            "geometry_type": self.geometry_type.value,
            "geometry_id": self.geometry_id,
            "origin": self.origin.tolist(),
            "direction": self.direction.tolist() if self.direction is not None else None,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> GeometryRef:
        """Construct from a dict produced by :meth:`to_dict`."""
        direction_raw = data.get("direction")
        return cls(
            body_id=data["body_id"],
            geometry_type=GeometryType(data["geometry_type"]),
            geometry_id=data["geometry_id"],
            origin=np.asarray(data["origin"], dtype=np.float64),
            direction=(
                np.asarray(direction_raw, dtype=np.float64) if direction_raw is not None else None
            ),
        )


@dataclass
class Mate:
    """A mate constraint between geometry on two bodies.

    Attributes:
        mate_id: Unique identifier for this mate.
        mate_type: The type of constraint (Coincident, Concentric, etc.).
        ref_a: Geometry reference on the first body.
        ref_b: Geometry reference on the second body.
        value: Scalar parameter for DISTANCE and ANGLE mates (0 otherwise).
        tolerance: Numeric tolerance for constraint satisfaction.
    """

    mate_id: int
    mate_type: MateType
    ref_a: GeometryRef
    ref_b: GeometryRef
    value: float = 0.0
    tolerance: float = 1e-6

    def validate(self) -> None:
        """Raise ``ValueError`` if this mate has incompatible geometry.

        Checks:
        - Self-mate (both refs on same body)
        - CONCENTRIC requires AXIS geometry on both refs
        - PARALLEL requires directional geometry (not POINT)
        - TANGENT requires surface geometry (FACE or EDGE)
        - Directional geometry types must have a direction vector
        """
        if self.ref_a.body_id == self.ref_b.body_id:
            msg = f"Self-mate: ref_a and ref_b both reference body {self.ref_a.body_id}"
            raise ValueError(msg)

        for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]:
            if ref.geometry_type in _DIRECTIONAL_TYPES and ref.direction is None:
                msg = (
                    f"{label}: geometry type {ref.geometry_type.value} requires a direction vector"
                )
                raise ValueError(msg)

        if self.mate_type is MateType.CONCENTRIC:
            for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]:
                if ref.geometry_type is not GeometryType.AXIS:
                    msg = (
                        f"CONCENTRIC mate requires AXIS geometry, "
                        f"got {ref.geometry_type.value} on {label}"
                    )
                    raise ValueError(msg)

        if self.mate_type is MateType.PARALLEL:
            for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]:
                if ref.geometry_type is GeometryType.POINT:
                    msg = f"PARALLEL mate requires directional geometry, got POINT on {label}"
                    raise ValueError(msg)

        if self.mate_type is MateType.TANGENT:
            _surface = frozenset({GeometryType.FACE, GeometryType.EDGE})
            for label, ref in [("ref_a", self.ref_a), ("ref_b", self.ref_b)]:
                if ref.geometry_type not in _surface:
                    msg = (
                        f"TANGENT mate requires surface geometry "
                        f"(FACE or EDGE), got {ref.geometry_type.value} "
                        f"on {label}"
                    )
                    raise ValueError(msg)

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable dict."""
        return {
            "mate_id": self.mate_id,
            "mate_type": self.mate_type.name,
            "ref_a": self.ref_a.to_dict(),
            "ref_b": self.ref_b.to_dict(),
            "value": self.value,
            "tolerance": self.tolerance,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Mate:
        """Construct from a dict produced by :meth:`to_dict`."""
        return cls(
            mate_id=data["mate_id"],
            mate_type=MateType[data["mate_type"]],
            ref_a=GeometryRef.from_dict(data["ref_a"]),
            ref_b=GeometryRef.from_dict(data["ref_b"]),
            value=data.get("value", 0.0),
            tolerance=data.get("tolerance", 1e-6),
        )


# ---------------------------------------------------------------------------
# Context-dependent DOF removal
# ---------------------------------------------------------------------------

# Lookup table: (MateType, ref_a GeometryType, ref_b GeometryType) -> DOF removed.
# Entries with None match any geometry type for that position.
_DOF_TABLE: dict[tuple[MateType, GeometryType | None, GeometryType | None], int] = {
    # COINCIDENT — context-dependent
    (MateType.COINCIDENT, GeometryType.FACE, GeometryType.FACE): 3,
    (MateType.COINCIDENT, GeometryType.POINT, GeometryType.POINT): 3,
    (MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE): 3,
    (MateType.COINCIDENT, GeometryType.EDGE, GeometryType.EDGE): 2,
    (MateType.COINCIDENT, GeometryType.FACE, GeometryType.POINT): 1,
    (MateType.COINCIDENT, GeometryType.POINT, GeometryType.FACE): 1,
    # CONCENTRIC
    (MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS): 2,
    # PARALLEL
    (MateType.PARALLEL, GeometryType.AXIS, GeometryType.AXIS): 2,
    (MateType.PARALLEL, GeometryType.FACE, GeometryType.FACE): 2,
    (MateType.PARALLEL, GeometryType.PLANE, GeometryType.PLANE): 2,
    # TANGENT
    (MateType.TANGENT, GeometryType.FACE, GeometryType.FACE): 1,
    (MateType.TANGENT, GeometryType.FACE, GeometryType.EDGE): 1,
    (MateType.TANGENT, GeometryType.EDGE, GeometryType.FACE): 1,
    # Types where DOF is always the same regardless of geometry
    (MateType.PERPENDICULAR, None, None): 1,
    (MateType.DISTANCE, None, None): 1,
    (MateType.ANGLE, None, None): 1,
    (MateType.LOCK, None, None): 6,
}


def dof_removed(
    mate_type: MateType,
    ref_a: GeometryRef,
    ref_b: GeometryRef,
) -> int:
    """Return the number of DOF removed by a mate given its geometry context.

    Looks up the exact ``(mate_type, ref_a.geometry_type, ref_b.geometry_type)``
    combination first, then falls back to a wildcard ``(mate_type, None, None)``
    entry, and finally to :attr:`MateType.default_dof`.

    Args:
        mate_type: The mate constraint type.
        ref_a: Geometry reference on the first body.
        ref_b: Geometry reference on the second body.

    Returns:
        Number of scalar DOF removed by this mate.
    """
    key = (mate_type, ref_a.geometry_type, ref_b.geometry_type)
    if key in _DOF_TABLE:
        return _DOF_TABLE[key]

    wildcard = (mate_type, None, None)
    if wildcard in _DOF_TABLE:
        return _DOF_TABLE[wildcard]

    return mate_type.default_dof
