"""Mate-based synthetic assembly generator.

Wraps SyntheticAssemblyGenerator to produce mate-level training data.
Generates joint-based assemblies via the existing generator, then
reverse-maps joints to plausible mate combinations. Supports noise
injection (redundant, missing, incompatible mates) for robust training.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np

from solver.datagen.generator import SyntheticAssemblyGenerator
from solver.datagen.types import Joint, JointType, RigidBody
from solver.mates.conversion import MateAnalysisResult, analyze_mate_assembly
from solver.mates.primitives import GeometryRef, GeometryType, Mate, MateType

if TYPE_CHECKING:
    from typing import Any

__all__ = [
    "SyntheticMateGenerator",
    "generate_mate_training_batch",
]


# ---------------------------------------------------------------------------
# Reverse mapping: JointType -> list of (MateType, geom_a, geom_b) combos
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class _MateSpec:
    """Specification for a mate to generate from a joint."""

    mate_type: MateType
    geom_a: GeometryType
    geom_b: GeometryType


_JOINT_TO_MATES: dict[JointType, list[_MateSpec]] = {
    JointType.REVOLUTE: [
        _MateSpec(MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS),
        _MateSpec(MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE),
    ],
    JointType.CYLINDRICAL: [
        _MateSpec(MateType.CONCENTRIC, GeometryType.AXIS, GeometryType.AXIS),
    ],
    JointType.BALL: [
        _MateSpec(MateType.COINCIDENT, GeometryType.POINT, GeometryType.POINT),
    ],
    JointType.FIXED: [
        _MateSpec(MateType.LOCK, GeometryType.FACE, GeometryType.FACE),
    ],
    JointType.SLIDER: [
        _MateSpec(MateType.COINCIDENT, GeometryType.PLANE, GeometryType.PLANE),
        _MateSpec(MateType.PARALLEL, GeometryType.AXIS, GeometryType.AXIS),
    ],
    JointType.PLANAR: [
        _MateSpec(MateType.COINCIDENT, GeometryType.FACE, GeometryType.FACE),
    ],
}


# ---------------------------------------------------------------------------
# Generator
# ---------------------------------------------------------------------------


class SyntheticMateGenerator:
    """Generates mate-based assemblies for training data.

    Wraps SyntheticAssemblyGenerator to produce joint-based assemblies,
    then reverse-maps each joint to a plausible set of mate constraints.

    Args:
        seed: Random seed for reproducibility.
        redundant_prob: Probability of injecting a redundant mate per joint.
        missing_prob: Probability of dropping a mate from a multi-mate pattern.
        incompatible_prob: Probability of injecting a mate with wrong geometry.
    """

    def __init__(
        self,
        seed: int = 42,
        *,
        redundant_prob: float = 0.0,
        missing_prob: float = 0.0,
        incompatible_prob: float = 0.0,
    ) -> None:
        self._joint_gen = SyntheticAssemblyGenerator(seed=seed)
        self._rng = np.random.default_rng(seed)
        self.redundant_prob = redundant_prob
        self.missing_prob = missing_prob
        self.incompatible_prob = incompatible_prob

    def _make_geometry_ref(
        self,
        body_id: int,
        geom_type: GeometryType,
        joint: Joint,
        *,
        is_ref_a: bool = True,
    ) -> GeometryRef:
        """Create a GeometryRef from joint geometry.

        Uses joint anchor, axis, and body_id to produce a ref
        with realistic geometry for the given type.
        """
        origin = joint.anchor_a if is_ref_a else joint.anchor_b

        direction: np.ndarray | None = None
        if geom_type in {GeometryType.AXIS, GeometryType.PLANE, GeometryType.FACE}:
            direction = joint.axis.copy()

        geom_id = f"{geom_type.value.capitalize()}001"

        return GeometryRef(
            body_id=body_id,
            geometry_type=geom_type,
            geometry_id=geom_id,
            origin=origin.copy(),
            direction=direction,
        )

    def _reverse_map_joint(
        self,
        joint: Joint,
        next_mate_id: int,
    ) -> list[Mate]:
        """Convert a joint to its mate representation."""
        specs = _JOINT_TO_MATES.get(joint.joint_type, [])
        if not specs:
            # Fallback: emit a single DISTANCE mate
            specs = [_MateSpec(MateType.DISTANCE, GeometryType.POINT, GeometryType.POINT)]

        mates: list[Mate] = []
        for spec in specs:
            ref_a = self._make_geometry_ref(joint.body_a, spec.geom_a, joint, is_ref_a=True)
            ref_b = self._make_geometry_ref(joint.body_b, spec.geom_b, joint, is_ref_a=False)
            mates.append(
                Mate(
                    mate_id=next_mate_id + len(mates),
                    mate_type=spec.mate_type,
                    ref_a=ref_a,
                    ref_b=ref_b,
                )
            )
        return mates

    def _inject_noise(
        self,
        mates: list[Mate],
        next_mate_id: int,
    ) -> list[Mate]:
        """Apply noise injection to the mate list.

        Modifies the list in-place and may add new mates.
        Returns the (possibly extended) list.
        """
        result = list(mates)
        extra: list[Mate] = []

        for mate in mates:
            # Redundant: duplicate a mate
            if self._rng.random() < self.redundant_prob:
                dup = Mate(
                    mate_id=next_mate_id + len(extra),
                    mate_type=mate.mate_type,
                    ref_a=mate.ref_a,
                    ref_b=mate.ref_b,
                    value=mate.value,
                    tolerance=mate.tolerance,
                )
                extra.append(dup)

            # Incompatible: wrong geometry type
            if self._rng.random() < self.incompatible_prob:
                bad_geom = GeometryType.POINT
                bad_ref = GeometryRef(
                    body_id=mate.ref_a.body_id,
                    geometry_type=bad_geom,
                    geometry_id="BadGeom001",
                    origin=mate.ref_a.origin.copy(),
                    direction=None,
                )
                extra.append(
                    Mate(
                        mate_id=next_mate_id + len(extra),
                        mate_type=MateType.CONCENTRIC,
                        ref_a=bad_ref,
                        ref_b=mate.ref_b,
                    )
                )

        result.extend(extra)

        # Missing: drop mates from multi-mate patterns (only if > 1 mate
        # for same body pair)
        if self.missing_prob > 0:
            filtered: list[Mate] = []
            for mate in result:
                if self._rng.random() < self.missing_prob:
                    continue
                filtered.append(mate)
            # Ensure at least one mate remains
            if not filtered and result:
                filtered = [result[0]]
            result = filtered

        return result

    def generate(
        self,
        n_bodies: int = 4,
        *,
        grounded: bool = False,
    ) -> tuple[list[RigidBody], list[Mate], MateAnalysisResult]:
        """Generate a mate-based assembly.

        Args:
            n_bodies: Number of rigid bodies.
            grounded: Whether to ground the first body.

        Returns:
            (bodies, mates, analysis_result) tuple.
        """
        bodies, joints, _analysis = self._joint_gen.generate_chain_assembly(
            n_bodies,
            joint_type=JointType.REVOLUTE,
            grounded=grounded,
        )

        mates: list[Mate] = []
        next_id = 0
        for joint in joints:
            joint_mates = self._reverse_map_joint(joint, next_id)
            mates.extend(joint_mates)
            next_id += len(joint_mates)

        # Apply noise
        mates = self._inject_noise(mates, next_id)

        ground_body = bodies[0].body_id if grounded else None
        result = analyze_mate_assembly(bodies, mates, ground_body)

        return bodies, mates, result


# ---------------------------------------------------------------------------
# Batch generation
# ---------------------------------------------------------------------------


def generate_mate_training_batch(
    batch_size: int = 100,
    n_bodies_range: tuple[int, int] = (3, 8),
    seed: int = 42,
    *,
    redundant_prob: float = 0.0,
    missing_prob: float = 0.0,
    incompatible_prob: float = 0.0,
    grounded_ratio: float = 1.0,
) -> list[dict[str, Any]]:
    """Produce a batch of mate-level training examples.

    Args:
        batch_size: Number of assemblies to generate.
        n_bodies_range: (min, max_exclusive) body count.
        seed: Random seed.
        redundant_prob: Probability of redundant mate injection.
        missing_prob: Probability of missing mate injection.
        incompatible_prob: Probability of incompatible mate injection.
        grounded_ratio: Fraction of assemblies that are grounded.

    Returns:
        List of dicts with bodies, mates, patterns, and labels.
    """
    rng = np.random.default_rng(seed)
    examples: list[dict[str, Any]] = []

    for i in range(batch_size):
        gen = SyntheticMateGenerator(
            seed=seed + i,
            redundant_prob=redundant_prob,
            missing_prob=missing_prob,
            incompatible_prob=incompatible_prob,
        )
        n = int(rng.integers(*n_bodies_range))
        grounded = bool(rng.random() < grounded_ratio)

        bodies, mates, result = gen.generate(n, grounded=grounded)

        examples.append(
            {
                "bodies": [
                    {
                        "body_id": b.body_id,
                        "position": b.position.tolist(),
                    }
                    for b in bodies
                ],
                "mates": [m.to_dict() for m in mates],
                "patterns": [p.to_dict() for p in result.patterns],
                "labels": result.labels.to_dict() if result.labels else None,
                "n_bodies": len(bodies),
                "n_mates": len(mates),
                "n_joints": len(result.joints),
            }
        )

    return examples
