"""Per-entity DOF diagnostics and overconstrained detection.

Provides per-body remaining degrees of freedom, human-readable free
motion labels, and redundant/conflicting constraint identification.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List

import numpy as np

from .entities import RigidBody
from .expr import Expr
from .params import ParamTable

# -- Per-entity DOF -----------------------------------------------------------


@dataclass
class EntityDOF:
    """DOF report for a single entity (rigid body)."""

    entity_id: str
    remaining_dof: int  # 0 = well-constrained
    free_motions: list[str] = field(default_factory=list)


def per_entity_dof(
    residuals: list[Expr],
    params: ParamTable,
    bodies: dict[str, RigidBody],
    rank_tol: float = 1e-8,
    jac_exprs: "list[list[Expr]] | None" = None,
) -> list[EntityDOF]:
    """Compute remaining DOF for each non-grounded body.

    For each body, extracts the Jacobian columns corresponding to its
    7 parameters, performs SVD to find constrained directions, and
    classifies null-space vectors as translations or rotations.
    """
    free = params.free_names()
    n_res = len(residuals)
    env = params.get_env()

    if n_res == 0:
        # No constraints — every free body has 6 DOF
        result = []
        for pid, body in bodies.items():
            if body.grounded:
                continue
            result.append(
                EntityDOF(
                    entity_id=pid,
                    remaining_dof=6,
                    free_motions=[
                        "translation along X",
                        "translation along Y",
                        "translation along Z",
                        "rotation about X",
                        "rotation about Y",
                        "rotation about Z",
                    ],
                )
            )
        return result

    # Build column index mapping: param_name -> column index in free list
    free_index = {name: i for i, name in enumerate(free)}

    # Build full Jacobian (for efficiency, compute once)
    n_free = len(free)
    J_full = np.empty((n_res, n_free))
    if jac_exprs is not None:
        for i in range(n_res):
            for j in range(n_free):
                J_full[i, j] = jac_exprs[i][j].eval(env)
    else:
        for i, r in enumerate(residuals):
            for j, name in enumerate(free):
                J_full[i, j] = r.diff(name).simplify().eval(env)

    result = []
    for pid, body in bodies.items():
        if body.grounded:
            continue

        # Find column indices for this body's params
        pfx = pid + "/"
        body_param_names = [
            pfx + "tx",
            pfx + "ty",
            pfx + "tz",
            pfx + "qw",
            pfx + "qx",
            pfx + "qy",
            pfx + "qz",
        ]
        col_indices = [free_index[n] for n in body_param_names if n in free_index]

        if not col_indices:
            # All params fixed (shouldn't happen for non-grounded, but be safe)
            result.append(EntityDOF(entity_id=pid, remaining_dof=0))
            continue

        # Extract submatrix: all residual rows, only this body's columns
        J_sub = J_full[:, col_indices]

        # SVD
        U, sv, Vt = np.linalg.svd(J_sub, full_matrices=True)
        constrained = int(np.sum(sv > rank_tol))

        # Subtract 1 for the quaternion unit-norm constraint (already in residuals)
        # The quat norm residual constrains 1 direction in the 7-D param space,
        # so effective body DOF = 7 - 1 - constrained_by_other_constraints.
        # But the quat norm IS one of the residual rows, so it's already counted
        # in `constrained`. So: remaining = len(col_indices) - constrained
        # But the quat norm takes 1 from 7 → 6 geometric DOF, and constrained
        # includes the quat norm row. So remaining = 7 - constrained, which gives
        # geometric remaining DOF directly.
        remaining = len(col_indices) - constrained

        # Classify null-space vectors as free motions
        free_motions = []
        if remaining > 0 and Vt.shape[0] > constrained:
            null_space = Vt[constrained:]  # rows = null vectors in param space

            # Map column indices back to param types
            param_types = []
            for n in body_param_names:
                if n in free_index:
                    if n.endswith(("/tx", "/ty", "/tz")):
                        param_types.append("t")
                    else:
                        param_types.append("q")

            for null_vec in null_space:
                label = _classify_motion(
                    null_vec, param_types, body_param_names, free_index
                )
                if label:
                    free_motions.append(label)

        result.append(
            EntityDOF(
                entity_id=pid,
                remaining_dof=remaining,
                free_motions=free_motions,
            )
        )

    return result


def _classify_motion(
    null_vec: np.ndarray,
    param_types: list[str],
    body_param_names: list[str],
    free_index: dict[str, int],
) -> str:
    """Classify a null-space vector as translation, rotation, or helical."""
    # Split components into translation and rotation parts
    trans_indices = [i for i, t in enumerate(param_types) if t == "t"]
    rot_indices = [i for i, t in enumerate(param_types) if t == "q"]

    trans_norm = np.linalg.norm(null_vec[trans_indices]) if trans_indices else 0.0
    rot_norm = np.linalg.norm(null_vec[rot_indices]) if rot_indices else 0.0

    total = trans_norm + rot_norm
    if total < 1e-14:
        return ""

    trans_frac = trans_norm / total
    rot_frac = rot_norm / total

    # Determine dominant axis
    if trans_frac > 0.8:
        # Pure translation
        axis = _dominant_axis(null_vec, trans_indices)
        return f"translation along {axis}"
    elif rot_frac > 0.8:
        # Pure rotation
        axis = _dominant_axis(null_vec, rot_indices)
        return f"rotation about {axis}"
    else:
        # Mixed — helical
        axis = _dominant_axis(null_vec, trans_indices)
        return f"helical motion along {axis}"


def _dominant_axis(vec: np.ndarray, indices: list[int]) -> str:
    """Find the dominant axis (X/Y/Z) among the given component indices."""
    if not indices:
        return "?"
    components = np.abs(vec[indices])
    # Map to axis names — first 3 in group are X/Y/Z
    axis_names = ["X", "Y", "Z"]
    if len(components) >= 3:
        idx = int(np.argmax(components[:3]))
        return axis_names[idx]
    elif len(components) == 1:
        return axis_names[0]
    else:
        idx = int(np.argmax(components))
        return axis_names[min(idx, 2)]


# -- Overconstrained detection ------------------------------------------------


@dataclass
class ConstraintDiag:
    """Diagnostic for a single constraint."""

    constraint_index: int
    kind: str  # "redundant" | "conflicting"
    detail: str


def find_overconstrained(
    residuals: list[Expr],
    params: ParamTable,
    residual_ranges: list[tuple[int, int, int]],
    rank_tol: float = 1e-8,
    jac_exprs: "list[list[Expr]] | None" = None,
) -> list[ConstraintDiag]:
    """Identify redundant and conflicting constraints.

    Algorithm (following SolvSpace's FindWhichToRemoveToFixJacobian):
    1. Build full Jacobian J, compute rank.
    2. If rank == n_residuals, not overconstrained — return empty.
    3. For each constraint: remove its rows, check if rank is preserved
       → if so, the constraint is **redundant**.
    4. Compute left null space, project residual vector F → if a
       constraint's residuals contribute to this projection, it is
       **conflicting** (redundant + unsatisfied).
    """
    free = params.free_names()
    n_free = len(free)
    n_res = len(residuals)

    if n_free == 0 or n_res == 0:
        return []

    env = params.get_env()

    # Build Jacobian and residual vector
    J = np.empty((n_res, n_free))
    r_vec = np.empty(n_res)
    for i, r in enumerate(residuals):
        r_vec[i] = r.eval(env)
    if jac_exprs is not None:
        for i in range(n_res):
            for j in range(n_free):
                J[i, j] = jac_exprs[i][j].eval(env)
    else:
        for i, r in enumerate(residuals):
            for j, name in enumerate(free):
                J[i, j] = r.diff(name).simplify().eval(env)

    # Full rank
    sv_full = np.linalg.svd(J, compute_uv=False)
    full_rank = int(np.sum(sv_full > rank_tol))

    if full_rank >= n_res:
        return []  # not overconstrained

    # Left null space: columns of U beyond rank
    U, sv, Vt = np.linalg.svd(J, full_matrices=True)
    left_null = U[:, full_rank:]  # shape (n_res, n_res - rank)

    # Project residual onto left null space
    null_residual = left_null.T @ r_vec  # shape (n_res - rank,)
    residual_projection = left_null @ null_residual  # back to residual space

    diags = []
    for start, end, c_idx in residual_ranges:
        # Remove this constraint's rows and check rank
        mask = np.ones(n_res, dtype=bool)
        mask[start:end] = False
        J_reduced = J[mask]

        if J_reduced.shape[0] == 0:
            continue

        sv_reduced = np.linalg.svd(J_reduced, compute_uv=False)
        reduced_rank = int(np.sum(sv_reduced > rank_tol))

        if reduced_rank >= full_rank:
            # Removing this constraint preserves rank → redundant
            # Check if it's also conflicting (contributes to unsatisfied null projection)
            constraint_proj = np.linalg.norm(residual_projection[start:end])
            if constraint_proj > rank_tol:
                kind = "conflicting"
                detail = (
                    f"Constraint {c_idx} is conflicting (redundant and unsatisfied)"
                )
            else:
                kind = "redundant"
                detail = (
                    f"Constraint {c_idx} is redundant (can be removed without effect)"
                )
            diags.append(
                ConstraintDiag(
                    constraint_index=c_idx,
                    kind=kind,
                    detail=detail,
                )
            )

    return diags
