"""Ground truth labeling pipeline for synthetic assemblies.

Produces rich per-constraint, per-joint, per-body, and assembly-level
labels by running both the pebble game and Jacobian verification and
correlating their results.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np

from solver.datagen.jacobian import JacobianVerifier
from solver.datagen.pebble_game import PebbleGame3D
from solver.datagen.types import (
    ConstraintAnalysis,
    Joint,
    JointType,
    RigidBody,
)

if TYPE_CHECKING:
    from typing import Any

__all__ = ["AssemblyLabels", "label_assembly"]

_GROUND_ID = -1
_SVD_TOL = 1e-8


# ---------------------------------------------------------------------------
# Label dataclasses
# ---------------------------------------------------------------------------


@dataclass
class ConstraintLabel:
    """Per scalar-constraint label combining both analysis methods."""

    joint_id: int
    constraint_idx: int
    pebble_independent: bool
    jacobian_independent: bool


@dataclass
class JointLabel:
    """Aggregated constraint counts for a single joint."""

    joint_id: int
    independent_count: int
    redundant_count: int
    total: int


@dataclass
class BodyDofLabel:
    """Per-body DOF signature from nullspace projection."""

    body_id: int
    translational_dof: int
    rotational_dof: int


@dataclass
class AssemblyLabel:
    """Assembly-wide summary label."""

    classification: str
    total_dof: int
    redundant_count: int
    is_rigid: bool
    is_minimally_rigid: bool
    has_degeneracy: bool


@dataclass
class AssemblyLabels:
    """Complete ground truth labels for an assembly."""

    per_constraint: list[ConstraintLabel]
    per_joint: list[JointLabel]
    per_body: list[BodyDofLabel]
    assembly: AssemblyLabel
    analysis: ConstraintAnalysis

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable dict."""
        return {
            "per_constraint": [
                {
                    "joint_id": c.joint_id,
                    "constraint_idx": c.constraint_idx,
                    "pebble_independent": c.pebble_independent,
                    "jacobian_independent": c.jacobian_independent,
                }
                for c in self.per_constraint
            ],
            "per_joint": [
                {
                    "joint_id": j.joint_id,
                    "independent_count": j.independent_count,
                    "redundant_count": j.redundant_count,
                    "total": j.total,
                }
                for j in self.per_joint
            ],
            "per_body": [
                {
                    "body_id": b.body_id,
                    "translational_dof": b.translational_dof,
                    "rotational_dof": b.rotational_dof,
                }
                for b in self.per_body
            ],
            "assembly": {
                "classification": self.assembly.classification,
                "total_dof": self.assembly.total_dof,
                "redundant_count": self.assembly.redundant_count,
                "is_rigid": self.assembly.is_rigid,
                "is_minimally_rigid": self.assembly.is_minimally_rigid,
                "has_degeneracy": self.assembly.has_degeneracy,
            },
        }


# ---------------------------------------------------------------------------
# Per-body DOF from nullspace projection
# ---------------------------------------------------------------------------


def _compute_per_body_dof(
    j_reduced: np.ndarray,
    body_ids: list[int],
    ground_body: int | None,
    body_index: dict[int, int],
) -> list[BodyDofLabel]:
    """Compute translational and rotational DOF per body.

    Uses SVD nullspace projection: for each body, extract its
    translational (3 cols) and rotational (3 cols) components
    from the nullspace basis and compute ranks.
    """
    # Build column index mapping for the reduced Jacobian
    # (ground body columns have been removed)
    col_map: dict[int, int] = {}
    col_idx = 0
    for bid in body_ids:
        if bid == ground_body:
            continue
        col_map[bid] = col_idx
        col_idx += 1

    results: list[BodyDofLabel] = []

    if j_reduced.size == 0:
        # No constraints — every body is fully free
        for bid in body_ids:
            if bid == ground_body:
                results.append(BodyDofLabel(body_id=bid, translational_dof=0, rotational_dof=0))
            else:
                results.append(BodyDofLabel(body_id=bid, translational_dof=3, rotational_dof=3))
        return results

    # Full SVD to get nullspace
    _u, s, vh = np.linalg.svd(j_reduced, full_matrices=True)
    rank = int(np.sum(s > _SVD_TOL))
    n_cols = j_reduced.shape[1]

    if rank >= n_cols:
        # Fully constrained — no nullspace
        for bid in body_ids:
            results.append(BodyDofLabel(body_id=bid, translational_dof=0, rotational_dof=0))
        return results

    # Nullspace basis: rows of Vh beyond the rank
    nullspace = vh[rank:]  # shape: (n_cols - rank, n_cols)

    for bid in body_ids:
        if bid == ground_body:
            results.append(BodyDofLabel(body_id=bid, translational_dof=0, rotational_dof=0))
            continue

        idx = col_map[bid]
        trans_cols = nullspace[:, idx * 6 : idx * 6 + 3]
        rot_cols = nullspace[:, idx * 6 + 3 : idx * 6 + 6]

        # Rank of each block = DOF in that category
        if trans_cols.size > 0:
            sv_t = np.linalg.svd(trans_cols, compute_uv=False)
            t_dof = int(np.sum(sv_t > _SVD_TOL))
        else:
            t_dof = 0

        if rot_cols.size > 0:
            sv_r = np.linalg.svd(rot_cols, compute_uv=False)
            r_dof = int(np.sum(sv_r > _SVD_TOL))
        else:
            r_dof = 0

        results.append(BodyDofLabel(body_id=bid, translational_dof=t_dof, rotational_dof=r_dof))

    return results


# ---------------------------------------------------------------------------
# Main labeling function
# ---------------------------------------------------------------------------


def label_assembly(
    bodies: list[RigidBody],
    joints: list[Joint],
    ground_body: int | None = None,
) -> AssemblyLabels:
    """Produce complete ground truth labels for an assembly.

    Runs both the pebble game and Jacobian verification internally,
    then correlates their results into per-constraint, per-joint,
    per-body, and assembly-level labels.

    Args:
        bodies: Rigid bodies in the assembly.
        joints: Joints connecting the bodies.
        ground_body: If set, this body is fixed to the world.

    Returns:
        AssemblyLabels with full label set and embedded ConstraintAnalysis.
    """
    # ---- Pebble Game ----
    pg = PebbleGame3D()
    all_edge_results: list[dict[str, Any]] = []

    if ground_body is not None:
        pg.add_body(_GROUND_ID)

    for body in bodies:
        pg.add_body(body.body_id)

    if ground_body is not None:
        ground_joint = Joint(
            joint_id=-1,
            body_a=ground_body,
            body_b=_GROUND_ID,
            joint_type=JointType.FIXED,
            anchor_a=bodies[0].position if bodies else np.zeros(3),
            anchor_b=bodies[0].position if bodies else np.zeros(3),
        )
        pg.add_joint(ground_joint)

    for joint in joints:
        results = pg.add_joint(joint)
        all_edge_results.extend(results)

    grounded = ground_body is not None
    combinatorial_independent = len(pg.state.independent_edges)
    raw_dof = pg.get_dof()
    ground_offset = 6 if grounded else 0
    effective_dof = raw_dof - ground_offset
    effective_internal_dof = max(0, effective_dof - (0 if grounded else 6))

    redundant_count = pg.get_redundant_count()
    if redundant_count > 0 and effective_internal_dof > 0:
        classification = "mixed"
    elif redundant_count > 0:
        classification = "overconstrained"
    elif effective_internal_dof > 0:
        classification = "underconstrained"
    else:
        classification = "well-constrained"

    # ---- Jacobian Verification ----
    verifier = JacobianVerifier(bodies)

    for joint in joints:
        verifier.add_joint_constraints(joint)

    j_full = verifier.get_jacobian()
    j_reduced = j_full.copy()
    if ground_body is not None and j_reduced.size > 0:
        idx = verifier.body_index[ground_body]
        cols_to_remove = list(range(idx * 6, (idx + 1) * 6))
        j_reduced = np.delete(j_reduced, cols_to_remove, axis=1)

    if j_reduced.size > 0:
        sv = np.linalg.svd(j_reduced, compute_uv=False)
        jacobian_rank = int(np.sum(sv > _SVD_TOL))
    else:
        jacobian_rank = 0

    n_cols = j_reduced.shape[1] if j_reduced.size > 0 else 6 * len(bodies)
    jacobian_nullity = n_cols - jacobian_rank
    dependent_rows = verifier.find_dependencies()
    dependent_set = set(dependent_rows)

    trivial_dof = 0 if grounded else 6
    jacobian_internal_dof = jacobian_nullity - trivial_dof
    geometric_degeneracies = max(0, combinatorial_independent - jacobian_rank)
    is_rigid = jacobian_nullity <= trivial_dof
    is_minimally_rigid = is_rigid and len(dependent_rows) == 0

    # ---- Per-constraint labels ----
    # Map Jacobian rows to (joint_id, constraint_index).
    # Rows are added contiguously per joint in the same order as joints.
    row_to_joint: list[tuple[int, int]] = []
    for joint in joints:
        dof = joint.joint_type.dof
        for ci in range(dof):
            row_to_joint.append((joint.joint_id, ci))

    per_constraint: list[ConstraintLabel] = []
    for edge_idx, edge_result in enumerate(all_edge_results):
        jid = edge_result["joint_id"]
        ci = edge_result["constraint_index"]
        pebble_indep = edge_result["independent"]

        # Find matching Jacobian row
        jacobian_indep = True
        if edge_idx < len(row_to_joint):
            row_idx = edge_idx
            jacobian_indep = row_idx not in dependent_set

        per_constraint.append(
            ConstraintLabel(
                joint_id=jid,
                constraint_idx=ci,
                pebble_independent=pebble_indep,
                jacobian_independent=jacobian_indep,
            )
        )

    # ---- Per-joint labels ----
    joint_agg: dict[int, JointLabel] = {}
    for cl in per_constraint:
        if cl.joint_id not in joint_agg:
            joint_agg[cl.joint_id] = JointLabel(
                joint_id=cl.joint_id,
                independent_count=0,
                redundant_count=0,
                total=0,
            )
        jl = joint_agg[cl.joint_id]
        jl.total += 1
        if cl.pebble_independent:
            jl.independent_count += 1
        else:
            jl.redundant_count += 1

    per_joint = [joint_agg[j.joint_id] for j in joints if j.joint_id in joint_agg]

    # ---- Per-body DOF labels ----
    body_ids = [b.body_id for b in bodies]
    per_body = _compute_per_body_dof(
        j_reduced,
        body_ids,
        ground_body,
        verifier.body_index,
    )

    # ---- Assembly label ----
    assembly_label = AssemblyLabel(
        classification=classification,
        total_dof=max(0, jacobian_internal_dof),
        redundant_count=redundant_count,
        is_rigid=is_rigid,
        is_minimally_rigid=is_minimally_rigid,
        has_degeneracy=geometric_degeneracies > 0,
    )

    # ---- ConstraintAnalysis (for backward compat) ----
    analysis = ConstraintAnalysis(
        combinatorial_dof=effective_dof,
        combinatorial_internal_dof=effective_internal_dof,
        combinatorial_redundant=redundant_count,
        combinatorial_classification=classification,
        per_edge_results=all_edge_results,
        jacobian_rank=jacobian_rank,
        jacobian_nullity=jacobian_nullity,
        jacobian_internal_dof=max(0, jacobian_internal_dof),
        numerically_dependent=dependent_rows,
        geometric_degeneracies=geometric_degeneracies,
        is_rigid=is_rigid,
        is_minimally_rigid=is_minimally_rigid,
    )

    return AssemblyLabels(
        per_constraint=per_constraint,
        per_joint=per_joint,
        per_body=per_body,
        assembly=assembly_label,
        analysis=analysis,
    )
