"""Numerical Jacobian rank verification for assembly constraint analysis.

Builds the constraint Jacobian matrix and analyzes its numerical rank
to detect geometric degeneracies that the combinatorial pebble game
cannot identify (e.g., parallel revolute axes creating hidden dependencies).

References:
    - Chappuis, "Constraints Derivation for Rigid Body Simulation in 3D"
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from solver.datagen.types import Joint, JointType, RigidBody

if TYPE_CHECKING:
    from typing import Any

__all__ = ["JacobianVerifier"]


class JacobianVerifier:
    """Builds and analyzes the constraint Jacobian for numerical rank check.

    The pebble game gives a combinatorial *necessary* condition for
    rigidity. However, geometric special cases (e.g., all revolute axes
    parallel, creating a hidden dependency) require numerical verification.

    For each joint, we construct the constraint Jacobian rows that map
    the 6n-dimensional generalized velocity vector to the constraint
    violation rates. The rank of this Jacobian equals the number of
    truly independent constraints.

    The generalized velocity vector for n bodies is::

        v = [v1_x, v1_y, v1_z, w1_x, w1_y, w1_z, ..., vn_x, ..., wn_z]

    Each scalar constraint C_i contributes one row to J such that::

        dC_i/dt = J_i @ v = 0
    """

    def __init__(self, bodies: list[RigidBody]) -> None:
        self.bodies = {b.body_id: b for b in bodies}
        self.body_index = {b.body_id: i for i, b in enumerate(bodies)}
        self.n_bodies = len(bodies)
        self.jacobian_rows: list[np.ndarray] = []
        self.row_labels: list[dict[str, Any]] = []

    def _body_cols(self, body_id: int) -> tuple[int, int]:
        """Return the column range [start, end) for a body in J."""
        idx = self.body_index[body_id]
        return idx * 6, (idx + 1) * 6

    def add_joint_constraints(self, joint: Joint) -> int:
        """Add Jacobian rows for all scalar constraints of a joint.

        Returns the number of rows added.
        """
        builder = {
            JointType.FIXED: self._build_fixed,
            JointType.REVOLUTE: self._build_revolute,
            JointType.CYLINDRICAL: self._build_cylindrical,
            JointType.SLIDER: self._build_slider,
            JointType.BALL: self._build_ball,
            JointType.PLANAR: self._build_planar,
            JointType.DISTANCE: self._build_distance,
            JointType.PARALLEL: self._build_parallel,
            JointType.PERPENDICULAR: self._build_perpendicular,
            JointType.UNIVERSAL: self._build_universal,
            JointType.SCREW: self._build_screw,
        }

        rows_before = len(self.jacobian_rows)
        builder[joint.joint_type](joint)
        return len(self.jacobian_rows) - rows_before

    def _make_row(self) -> np.ndarray:
        """Create a zero row of width 6*n_bodies."""
        return np.zeros(6 * self.n_bodies)

    def _skew(self, v: np.ndarray) -> np.ndarray:
        """Skew-symmetric matrix for cross product: ``skew(v) @ w = v x w``."""
        return np.array(
            [
                [0, -v[2], v[1]],
                [v[2], 0, -v[0]],
                [-v[1], v[0], 0],
            ]
        )

    # --- Ball-and-socket (spherical) joint: 3 translation constraints ---

    def _build_ball(self, joint: Joint) -> None:
        """Ball joint: coincident point constraint.

        ``C_trans = (x_b + R_b @ r_b) - (x_a + R_a @ r_a) = 0``
        (3 equations)

        Jacobian rows (for each of x, y, z):
            body_a linear:  -I
            body_a angular: +skew(R_a @ r_a)
            body_b linear:  +I
            body_b angular: -skew(R_b @ r_b)
        """
        # Use anchor positions directly as world-frame offsets
        r_a = joint.anchor_a - self.bodies[joint.body_a].position
        r_b = joint.anchor_b - self.bodies[joint.body_b].position

        col_a_start, col_a_end = self._body_cols(joint.body_a)
        col_b_start, col_b_end = self._body_cols(joint.body_b)

        for axis_idx in range(3):
            row = self._make_row()
            e = np.zeros(3)
            e[axis_idx] = 1.0

            row[col_a_start : col_a_start + 3] = -e
            row[col_a_start + 3 : col_a_end] = np.cross(r_a, e)

            row[col_b_start : col_b_start + 3] = e
            row[col_b_start + 3 : col_b_end] = -np.cross(r_b, e)

            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "ball_translation",
                    "axis": axis_idx,
                }
            )

    # --- Fixed joint: 3 translation + 3 rotation constraints ---

    def _build_fixed(self, joint: Joint) -> None:
        """Fixed joint = ball joint + locked rotation.

        Translation part: same as ball joint (3 rows).
        Rotation part: relative angular velocity must be zero (3 rows).
        """
        self._build_ball(joint)

        col_a_start, _ = self._body_cols(joint.body_a)
        col_b_start, _ = self._body_cols(joint.body_b)

        for axis_idx in range(3):
            row = self._make_row()
            row[col_a_start + 3 + axis_idx] = -1.0
            row[col_b_start + 3 + axis_idx] = 1.0

            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "fixed_rotation",
                    "axis": axis_idx,
                }
            )

    # --- Revolute (hinge) joint: 3 translation + 2 rotation constraints ---

    def _build_revolute(self, joint: Joint) -> None:
        """Revolute joint: rotation only about one axis.

        Translation: same as ball (3 rows).
        Rotation: relative angular velocity must be parallel to hinge axis.
        """
        self._build_ball(joint)

        axis = joint.axis / np.linalg.norm(joint.axis)
        t1, t2 = self._perpendicular_pair(axis)

        col_a_start, _ = self._body_cols(joint.body_a)
        col_b_start, _ = self._body_cols(joint.body_b)

        for i, t in enumerate((t1, t2)):
            row = self._make_row()
            row[col_a_start + 3 : col_a_start + 6] = -t
            row[col_b_start + 3 : col_b_start + 6] = t

            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "revolute_rotation",
                    "perp_axis": i,
                }
            )

    # --- Cylindrical joint: 2 translation + 2 rotation constraints ---

    def _build_cylindrical(self, joint: Joint) -> None:
        """Cylindrical joint: allows rotation + translation along one axis.

        Translation: constrain motion perpendicular to axis (2 rows).
        Rotation: constrain rotation perpendicular to axis (2 rows).
        """
        axis = joint.axis / np.linalg.norm(joint.axis)
        t1, t2 = self._perpendicular_pair(axis)

        r_a = joint.anchor_a - self.bodies[joint.body_a].position
        r_b = joint.anchor_b - self.bodies[joint.body_b].position

        col_a_start, col_a_end = self._body_cols(joint.body_a)
        col_b_start, col_b_end = self._body_cols(joint.body_b)

        for i, t in enumerate((t1, t2)):
            row = self._make_row()
            row[col_a_start : col_a_start + 3] = -t
            row[col_a_start + 3 : col_a_end] = np.cross(r_a, t)
            row[col_b_start : col_b_start + 3] = t
            row[col_b_start + 3 : col_b_end] = -np.cross(r_b, t)

            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "cylindrical_translation",
                    "perp_axis": i,
                }
            )

        for i, t in enumerate((t1, t2)):
            row = self._make_row()
            row[col_a_start + 3 : col_a_start + 6] = -t
            row[col_b_start + 3 : col_b_start + 6] = t

            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "cylindrical_rotation",
                    "perp_axis": i,
                }
            )

    # --- Slider (prismatic) joint: 2 translation + 3 rotation constraints ---

    def _build_slider(self, joint: Joint) -> None:
        """Slider/prismatic joint: translation along one axis only.

        Translation: perpendicular translation constrained (2 rows).
        Rotation: all relative rotation constrained (3 rows).
        """
        axis = joint.axis / np.linalg.norm(joint.axis)
        t1, t2 = self._perpendicular_pair(axis)

        r_a = joint.anchor_a - self.bodies[joint.body_a].position
        r_b = joint.anchor_b - self.bodies[joint.body_b].position

        col_a_start, col_a_end = self._body_cols(joint.body_a)
        col_b_start, col_b_end = self._body_cols(joint.body_b)

        for i, t in enumerate((t1, t2)):
            row = self._make_row()
            row[col_a_start : col_a_start + 3] = -t
            row[col_a_start + 3 : col_a_end] = np.cross(r_a, t)
            row[col_b_start : col_b_start + 3] = t
            row[col_b_start + 3 : col_b_end] = -np.cross(r_b, t)

            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "slider_translation",
                    "perp_axis": i,
                }
            )

        for axis_idx in range(3):
            row = self._make_row()
            row[col_a_start + 3 + axis_idx] = -1.0
            row[col_b_start + 3 + axis_idx] = 1.0

            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "slider_rotation",
                    "axis": axis_idx,
                }
            )

    # --- Planar joint: 1 translation + 2 rotation constraints ---

    def _build_planar(self, joint: Joint) -> None:
        """Planar joint: constrains to a plane.

        Translation: motion along plane normal constrained (1 row).
        Rotation: rotation about axes in the plane constrained (2 rows).
        """
        normal = joint.axis / np.linalg.norm(joint.axis)
        t1, t2 = self._perpendicular_pair(normal)

        r_a = joint.anchor_a - self.bodies[joint.body_a].position
        r_b = joint.anchor_b - self.bodies[joint.body_b].position

        col_a_start, col_a_end = self._body_cols(joint.body_a)
        col_b_start, col_b_end = self._body_cols(joint.body_b)

        row = self._make_row()
        row[col_a_start : col_a_start + 3] = -normal
        row[col_a_start + 3 : col_a_end] = np.cross(r_a, normal)
        row[col_b_start : col_b_start + 3] = normal
        row[col_b_start + 3 : col_b_end] = -np.cross(r_b, normal)
        self.jacobian_rows.append(row)
        self.row_labels.append(
            {
                "joint_id": joint.joint_id,
                "type": "planar_translation",
            }
        )

        for i, t in enumerate((t1, t2)):
            row = self._make_row()
            row[col_a_start + 3 : col_a_start + 6] = -t
            row[col_b_start + 3 : col_b_start + 6] = t
            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "planar_rotation",
                    "perp_axis": i,
                }
            )

    # --- Distance constraint: 1 scalar ---

    def _build_distance(self, joint: Joint) -> None:
        """Distance constraint: ``||p_b - p_a|| = d``.

        Single row: ``direction . (v_b + w_b x r_b - v_a - w_a x r_a) = 0``
        where ``direction = normalized(p_b - p_a)``.
        """
        p_a = joint.anchor_a
        p_b = joint.anchor_b
        diff = p_b - p_a
        dist = np.linalg.norm(diff)
        direction = np.array([1.0, 0.0, 0.0]) if dist < 1e-12 else diff / dist

        r_a = joint.anchor_a - self.bodies[joint.body_a].position
        r_b = joint.anchor_b - self.bodies[joint.body_b].position

        col_a_start, col_a_end = self._body_cols(joint.body_a)
        col_b_start, col_b_end = self._body_cols(joint.body_b)

        row = self._make_row()
        row[col_a_start : col_a_start + 3] = -direction
        row[col_a_start + 3 : col_a_end] = np.cross(r_a, direction)
        row[col_b_start : col_b_start + 3] = direction
        row[col_b_start + 3 : col_b_end] = -np.cross(r_b, direction)

        self.jacobian_rows.append(row)
        self.row_labels.append(
            {
                "joint_id": joint.joint_id,
                "type": "distance",
            }
        )

    # --- Parallel constraint: 3 rotation constraints ---

    def _build_parallel(self, joint: Joint) -> None:
        """Parallel: all relative rotation constrained (same as fixed rotation).

        In practice only 2 of 3 are independent for a single axis, but
        we emit 3 and let the rank check sort it out.
        """
        col_a_start, _ = self._body_cols(joint.body_a)
        col_b_start, _ = self._body_cols(joint.body_b)

        for axis_idx in range(3):
            row = self._make_row()
            row[col_a_start + 3 + axis_idx] = -1.0
            row[col_b_start + 3 + axis_idx] = 1.0
            self.jacobian_rows.append(row)
            self.row_labels.append(
                {
                    "joint_id": joint.joint_id,
                    "type": "parallel_rotation",
                    "axis": axis_idx,
                }
            )

    # --- Perpendicular constraint: 1 angular ---

    def _build_perpendicular(self, joint: Joint) -> None:
        """Perpendicular: single dot-product angular constraint."""
        axis = joint.axis / np.linalg.norm(joint.axis)
        col_a_start, _ = self._body_cols(joint.body_a)
        col_b_start, _ = self._body_cols(joint.body_b)

        row = self._make_row()
        row[col_a_start + 3 : col_a_start + 6] = -axis
        row[col_b_start + 3 : col_b_start + 6] = axis
        self.jacobian_rows.append(row)
        self.row_labels.append(
            {
                "joint_id": joint.joint_id,
                "type": "perpendicular",
            }
        )

    # --- Universal (Cardan) joint: 3 translation + 1 rotation ---

    def _build_universal(self, joint: Joint) -> None:
        """Universal joint: ball + one rotation constraint.

        Allows rotation about two axes, constrains rotation about the third.
        """
        self._build_ball(joint)

        axis = joint.axis / np.linalg.norm(joint.axis)
        col_a_start, _ = self._body_cols(joint.body_a)
        col_b_start, _ = self._body_cols(joint.body_b)

        row = self._make_row()
        row[col_a_start + 3 : col_a_start + 6] = -axis
        row[col_b_start + 3 : col_b_start + 6] = axis
        self.jacobian_rows.append(row)
        self.row_labels.append(
            {
                "joint_id": joint.joint_id,
                "type": "universal_rotation",
            }
        )

    # --- Screw (helical) joint: 2 translation + 2 rotation + 1 coupled ---

    def _build_screw(self, joint: Joint) -> None:
        """Screw joint: coupled rotation-translation along axis.

        Like cylindrical but with a coupling constraint:
        ``v_axial - pitch * w_axial = 0``
        """
        self._build_cylindrical(joint)

        axis = joint.axis / np.linalg.norm(joint.axis)
        col_a_start, _ = self._body_cols(joint.body_a)
        col_b_start, _ = self._body_cols(joint.body_b)

        row = self._make_row()
        row[col_a_start : col_a_start + 3] = -axis
        row[col_b_start : col_b_start + 3] = axis
        row[col_a_start + 3 : col_a_start + 6] = joint.pitch * axis
        row[col_b_start + 3 : col_b_start + 6] = -joint.pitch * axis
        self.jacobian_rows.append(row)
        self.row_labels.append(
            {
                "joint_id": joint.joint_id,
                "type": "screw_coupling",
            }
        )

    # --- Utilities ---

    def _perpendicular_pair(self, axis: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Generate two unit vectors perpendicular to *axis* and each other."""
        if abs(axis[0]) < 0.9:
            t1 = np.cross(axis, np.array([1.0, 0, 0]))
        else:
            t1 = np.cross(axis, np.array([0, 1.0, 0]))
        t1 /= np.linalg.norm(t1)
        t2 = np.cross(axis, t1)
        t2 /= np.linalg.norm(t2)
        return t1, t2

    def get_jacobian(self) -> np.ndarray:
        """Return the full constraint Jacobian matrix."""
        if not self.jacobian_rows:
            return np.zeros((0, 6 * self.n_bodies))
        return np.array(self.jacobian_rows)

    def numerical_rank(self, tol: float = 1e-8) -> int:
        """Compute the numerical rank of the constraint Jacobian via SVD.

        This is the number of truly independent scalar constraints,
        accounting for geometric degeneracies that the combinatorial
        pebble game cannot detect.
        """
        j = self.get_jacobian()
        if j.size == 0:
            return 0
        sv = np.linalg.svd(j, compute_uv=False)
        return int(np.sum(sv > tol))

    def find_dependencies(self, tol: float = 1e-8) -> list[int]:
        """Identify which constraint rows are numerically dependent.

        Returns indices of rows that can be removed without changing
        the Jacobian's rank.
        """
        j = self.get_jacobian()
        if j.size == 0:
            return []

        n_rows = j.shape[0]
        dependent: list[int] = []

        current = np.zeros((0, j.shape[1]))
        current_rank = 0

        for i in range(n_rows):
            candidate = np.vstack([current, j[i : i + 1, :]]) if current.size else j[i : i + 1, :]
            sv = np.linalg.svd(candidate, compute_uv=False)
            new_rank = int(np.sum(sv > tol))

            if new_rank > current_rank:
                current = candidate
                current_rank = new_rank
            else:
                dependent.append(i)

        return dependent
