"""Combined pebble game + Jacobian verification analysis.

Provides :func:`analyze_assembly`, the main entry point for full rigidity
analysis of an assembly using both combinatorial and numerical methods.
"""

from __future__ import annotations

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,
)

__all__ = ["analyze_assembly"]

_GROUND_ID = -1


def analyze_assembly(
    bodies: list[RigidBody],
    joints: list[Joint],
    ground_body: int | None = None,
) -> ConstraintAnalysis:
    """Full rigidity analysis of an assembly using both methods.

    Args:
        bodies: List of rigid bodies in the assembly.
        joints: List of joints connecting bodies.
        ground_body: If set, this body is fixed (adds 6 implicit constraints).

    Returns:
        ConstraintAnalysis with combinatorial and numerical results.
    """
    # --- Pebble Game ---
    pg = PebbleGame3D()
    all_edge_results = []

    # Add a virtual ground body (id=-1) if grounding is requested.
    # Grounding body X means adding a fixed joint between X and
    # the virtual ground. This properly lets the pebble game account
    # for the 6 removed DOF without breaking invariants.
    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)
        # Don't include ground joint edges in the output labels
        # (they're infrastructure, not user constraints)

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

    combinatorial_independent = len(pg.state.independent_edges)
    grounded = ground_body is not None

    # The virtual ground body contributes 6 pebbles to the total.
    # Subtract those from the reported DOF for user-facing numbers.
    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))

    # Classify based on effective (adjusted) DOF, not raw pebble game output,
    # because the virtual ground body skews the raw numbers.
    redundant = pg.get_redundant_count()
    if redundant > 0 and effective_internal_dof > 0:
        combinatorial_classification = "mixed"
    elif redundant > 0:
        combinatorial_classification = "overconstrained"
    elif effective_internal_dof > 0:
        combinatorial_classification = "underconstrained"
    else:
        combinatorial_classification = "well-constrained"

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

    for joint in joints:
        verifier.add_joint_constraints(joint)

    # If grounded, remove the ground body's columns (fix its DOF)
    j = verifier.get_jacobian()
    if ground_body is not None and j.size > 0:
        idx = verifier.body_index[ground_body]
        cols_to_remove = list(range(idx * 6, (idx + 1) * 6))
        j = np.delete(j, cols_to_remove, axis=1)

    if j.size > 0:
        sv = np.linalg.svd(j, compute_uv=False)
        jacobian_rank = int(np.sum(sv > 1e-8))
    else:
        jacobian_rank = 0

    n_cols = j.shape[1] if j.size > 0 else 6 * len(bodies)
    jacobian_nullity = n_cols - jacobian_rank

    dependent = verifier.find_dependencies()

    # Adjust for ground
    trivial_dof = 0 if ground_body is not None else 6
    jacobian_internal_dof = jacobian_nullity - trivial_dof

    geometric_degeneracies = max(0, combinatorial_independent - jacobian_rank)

    # Rigidity: numerically rigid if nullity == trivial DOF
    is_rigid = jacobian_nullity <= trivial_dof
    is_minimally_rigid = is_rigid and len(dependent) == 0

    return ConstraintAnalysis(
        combinatorial_dof=effective_dof,
        combinatorial_internal_dof=effective_internal_dof,
        combinatorial_redundant=pg.get_redundant_count(),
        combinatorial_classification=combinatorial_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,
        geometric_degeneracies=geometric_degeneracies,
        is_rigid=is_rigid,
        is_minimally_rigid=is_minimally_rigid,
    )
