"""Constraint classes that produce residual Expr lists.

Each constraint takes two RigidBody entities and marker transforms,
then generates residual expressions that equal zero when satisfied.

Phase 1 constraints: Coincident, DistancePointPoint, Fixed
Phase 2 constraints: all remaining BaseJointKind types from Types.h
"""

from __future__ import annotations

import math
from typing import List

from .entities import RigidBody
from .expr import Const, Expr
from .geometry import (
    cross3,
    dot3,
    marker_x_axis,
    marker_y_axis,
    marker_z_axis,
    point_line_perp_components,
    point_plane_distance,
    sub3,
)


class ConstraintBase:
    """Abstract constraint producing residual expressions."""

    def residuals(self) -> List[Expr]:
        raise NotImplementedError


class CoincidentConstraint(ConstraintBase):
    """Point-on-point: marker origin on body_i == marker origin on body_j.

    3 residuals: dx, dy, dz.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_j_pos = marker_j_pos

    def residuals(self) -> List[Expr]:
        wx_i, wy_i, wz_i = self.body_i.world_point(*self.marker_i_pos)
        wx_j, wy_j, wz_j = self.body_j.world_point(*self.marker_j_pos)
        return [wx_i - wx_j, wy_i - wy_j, wz_i - wz_j]


class DistancePointPointConstraint(ConstraintBase):
    """Point-to-point distance: ||p_i - p_j||^2 = d^2.

    1 residual (squared form avoids sqrt in Jacobian).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        distance: float,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_j_pos = marker_j_pos
        self.distance = distance

    def world_points(self) -> tuple[tuple[Expr, Expr, Expr], tuple[Expr, Expr, Expr]]:
        """Return (world_point_i, world_point_j) expression tuples."""
        return (
            self.body_i.world_point(*self.marker_i_pos),
            self.body_j.world_point(*self.marker_j_pos),
        )

    def residuals(self) -> List[Expr]:
        (wx_i, wy_i, wz_i), (wx_j, wy_j, wz_j) = self.world_points()
        dx = wx_i - wx_j
        dy = wy_i - wy_j
        dz = wz_i - wz_j
        d2 = dx * dx + dy * dy + dz * dz
        return [d2 - Const(self.distance * self.distance)]


class FixedConstraint(ConstraintBase):
    """Lock all 6 DOF of body_j relative to body_i.

    Position: 3 residuals (marker origins coincide).
    Orientation: 3 residuals from relative quaternion (imaginary part == 0
    when orientations match).  The quaternion normalization constraint
    handles the 4th equation.

    The reference relative quaternion q_ref = conj(q_i_marker) * q_j_marker
    is computed from the marker quaternions at constraint creation time.
    At solve time:  q_rel = conj(q_i_world * q_i_marker) * (q_j_world * q_j_marker)
    Residual = imaginary part of (conj(q_ref) * q_rel).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        # Position residuals — same as Coincident
        wx_i, wy_i, wz_i = self.body_i.world_point(*self.marker_i_pos)
        wx_j, wy_j, wz_j = self.body_j.world_point(*self.marker_j_pos)
        pos_res = [wx_i - wx_j, wy_i - wy_j, wz_i - wz_j]

        # Orientation residuals via quaternion error
        # q_i_total = q_body_i * q_marker_i  (as Expr)
        # q_j_total = q_body_j * q_marker_j
        # q_error = conj(q_i_total) * q_j_total
        # We want q_error ≈ identity → imaginary parts ≈ 0
        qi = _quat_mul_const(
            self.body_i.qw,
            self.body_i.qx,
            self.body_i.qy,
            self.body_i.qz,
            *self.marker_i_quat,
        )
        qj = _quat_mul_const(
            self.body_j.qw,
            self.body_j.qx,
            self.body_j.qy,
            self.body_j.qz,
            *self.marker_j_quat,
        )
        # conj(qi) * qj
        err = _quat_mul_expr(
            qi[0],
            -qi[1],
            -qi[2],
            -qi[3],  # conjugate
            qj[0],
            qj[1],
            qj[2],
            qj[3],
        )
        # err should be (1, 0, 0, 0) — residuals are the imaginary parts
        ori_res = [err[1], err[2], err[3]]

        return pos_res + ori_res


# ============================================================================
# Phase 2: Point constraints
# ============================================================================


class PointOnLineConstraint(ConstraintBase):
    """Point constrained to a line — 2 DOF removed.

    marker_i origin lies on the line through marker_j origin along
    marker_j Z-axis.  2 residuals: perpendicular distance components.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = point_line_perp_components(p_i, p_j, z_j)
        return [cx, cy, cz]


class PointInPlaneConstraint(ConstraintBase):
    """Point constrained to a plane — 1 DOF removed.

    marker_i origin lies in the plane through marker_j origin with
    normal = marker_j Z-axis.  Optional offset via params[0].
    1 residual: signed distance to plane.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
        offset: float = 0.0,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat
        self.offset = offset

    def residuals(self) -> List[Expr]:
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        n_j = marker_z_axis(self.body_j, self.marker_j_quat)
        d = point_plane_distance(p_i, p_j, n_j)
        if self.offset != 0.0:
            d = d - Const(self.offset)
        return [d]


# ============================================================================
# Phase 2: Axis orientation constraints
# ============================================================================


class ParallelConstraint(ConstraintBase):
    """Parallel axes — 2 DOF removed.

    marker Z-axes are parallel: z_i x z_j = 0.
    3 cross-product residuals (rank 2 at the solution).  Using all 3
    avoids a singularity when the axes lie in the XY plane, where
    dropping cz would leave the constraint blind to yaw rotations.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_quat = marker_i_quat
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = cross3(z_i, z_j)
        return [cx, cy, cz]


class PerpendicularConstraint(ConstraintBase):
    """Perpendicular axes — 1 DOF removed.

    marker Z-axes are perpendicular: z_i . z_j = 0.
    1 residual.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_quat = marker_i_quat
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        return [dot3(z_i, z_j)]


class AngleConstraint(ConstraintBase):
    """Angle between axes — 1 DOF removed.

    z_i . z_j = cos(angle).
    1 residual.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_quat: tuple[float, float, float, float],
        angle: float,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_quat = marker_i_quat
        self.marker_j_quat = marker_j_quat
        self.angle = angle

    def residuals(self) -> List[Expr]:
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        return [dot3(z_i, z_j) - Const(math.cos(self.angle))]


# ============================================================================
# Phase 2: Axis/surface constraints
# ============================================================================


class ConcentricConstraint(ConstraintBase):
    """Coaxial / concentric — 4 DOF removed.

    Axes are collinear: parallel Z-axes (2) + point-on-line (2).
    Optional distance offset along axis via params.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
        distance: float = 0.0,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat
        self.distance = distance

    def residuals(self) -> List[Expr]:
        # Parallel axes (3 cross-product residuals, rank 2 at solution)
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = cross3(z_i, z_j)

        # Point-on-line: marker_i origin on line through marker_j along z_j
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        lx, ly, lz = point_line_perp_components(p_i, p_j, z_j)

        return [cx, cy, cz, lx, ly, lz]


class TangentConstraint(ConstraintBase):
    """Face-on-face tangency — 1 DOF removed.

    Signed distance between marker origins along marker_j normal = 0.
    1 residual: (p_i - p_j) . z_j
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        n_j = marker_z_axis(self.body_j, self.marker_j_quat)
        return [point_plane_distance(p_i, p_j, n_j)]


class PlanarConstraint(ConstraintBase):
    """Coplanar faces — 3 DOF removed.

    Parallel normals (2) + point-in-plane (1).  Optional offset.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
        offset: float = 0.0,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat
        self.offset = offset

    def residuals(self) -> List[Expr]:
        # Parallel normals (3 cross-product residuals, rank 2 at solution)
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = cross3(z_i, z_j)

        # Point-in-plane
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        d = point_plane_distance(p_i, p_j, z_j)
        if self.offset != 0.0:
            d = d - Const(self.offset)

        return [cx, cy, cz, d]


class LineInPlaneConstraint(ConstraintBase):
    """Line constrained to a plane — 2 DOF removed.

    Line defined by marker_i Z-axis lies in plane defined by marker_j normal.
    2 residuals: point-in-plane (1) + line direction perpendicular to normal (1).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
        offset: float = 0.0,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat
        self.offset = offset

    def residuals(self) -> List[Expr]:
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        n_j = marker_z_axis(self.body_j, self.marker_j_quat)
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)

        # Point in plane
        d = point_plane_distance(p_i, p_j, n_j)
        if self.offset != 0.0:
            d = d - Const(self.offset)

        # Line direction perpendicular to plane normal
        dir_dot = dot3(z_i, n_j)

        return [d, dir_dot]


# ============================================================================
# Phase 2: Kinematic joints
# ============================================================================


class BallConstraint(ConstraintBase):
    """Spherical joint — 3 DOF removed.

    Coincident marker origins.  Same as CoincidentConstraint.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
    ):
        self._inner = CoincidentConstraint(body_i, marker_i_pos, body_j, marker_j_pos)

    def residuals(self) -> List[Expr]:
        return self._inner.residuals()


class RevoluteConstraint(ConstraintBase):
    """Hinge joint — 5 DOF removed.

    Coincident origins (3) + parallel Z-axes (2).
    1 rotational DOF remains (about the common Z-axis).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        # Coincident origins
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        pos = [p_i[0] - p_j[0], p_i[1] - p_j[1], p_i[2] - p_j[2]]

        # Parallel Z-axes (3 cross-product residuals, rank 2 at solution)
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = cross3(z_i, z_j)

        return pos + [cx, cy, cz]


class CylindricalConstraint(ConstraintBase):
    """Cylindrical joint — 4 DOF removed.

    Parallel Z-axes (2) + point-on-line (2).
    2 DOF remain: rotation about and translation along the common axis.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        # Parallel Z-axes (3 cross-product residuals, rank 2 at solution)
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = cross3(z_i, z_j)

        # Point-on-line
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        lx, ly, lz = point_line_perp_components(p_i, p_j, z_j)

        return [cx, cy, cz, lx, ly, lz]


class SliderConstraint(ConstraintBase):
    """Prismatic / slider joint — 5 DOF removed.

    Parallel Z-axes (2) + point-on-line (2) + rotation lock (1).
    1 DOF remains: translation along the common Z-axis.

    Rotation lock: x_i . y_j = 0  (prevents twist about Z).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        # Parallel Z-axes (3 cross-product residuals, rank 2 at solution)
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = cross3(z_i, z_j)

        # Point-on-line
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        lx, ly, lz = point_line_perp_components(p_i, p_j, z_j)

        # Rotation lock: x_i . y_j = 0
        x_i = marker_x_axis(self.body_i, self.marker_i_quat)
        y_j = marker_y_axis(self.body_j, self.marker_j_quat)
        twist = dot3(x_i, y_j)

        return [cx, cy, cz, lx, ly, lz, twist]


class ScrewConstraint(ConstraintBase):
    """Helical / screw joint — 5 DOF removed.

    Cylindrical (4) + coupled rotation-translation via pitch (1).
    1 DOF remains: screw motion (rotation + proportional translation).

    The coupling residual uses the relative quaternion's Z-component
    (proportional to the rotation angle for small angles) and the axial
    displacement: axial_disp - pitch * (2 * qz_rel / qw_rel) / (2*pi) = 0.
    For the Newton solver operating near the solution, the linear
    approximation angle ≈ 2 * qz_rel is adequate.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
        pitch: float = 1.0,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat
        self.pitch = pitch

    def residuals(self) -> List[Expr]:
        # Cylindrical residuals (5: 3 parallel + 2 point-on-line)
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        cx, cy, cz = cross3(z_i, z_j)

        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        lx, ly, lz = point_line_perp_components(p_i, p_j, z_j)

        # Pitch coupling: axial_disp = pitch * angle / (2*pi)
        # Axial displacement
        d = sub3(p_i, p_j)
        axial = dot3(d, z_j)

        # Relative rotation about Z via quaternion
        # q_rel = conj(q_i_total) * q_j_total
        qi = _quat_mul_const(
            self.body_i.qw,
            self.body_i.qx,
            self.body_i.qy,
            self.body_i.qz,
            *self.marker_i_quat,
        )
        qj = _quat_mul_const(
            self.body_j.qw,
            self.body_j.qx,
            self.body_j.qy,
            self.body_j.qz,
            *self.marker_j_quat,
        )
        rel = _quat_mul_expr(qi[0], -qi[1], -qi[2], -qi[3], qj[0], qj[1], qj[2], qj[3])
        # For small angles: angle ≈ 2 * qz_rel, but qw_rel ≈ 1
        # Use sin(angle/2) form: residual = axial - pitch * 2*qz / (2*pi)
        # = axial - pitch * qz / pi
        coupling = axial - Const(self.pitch / math.pi) * rel[3]

        return [cx, cy, cz, lx, ly, lz, coupling]


class UniversalConstraint(ConstraintBase):
    """Universal / Cardan joint — 4 DOF removed.

    Coincident origins (3) + perpendicular Z-axes (1).
    2 DOF remain: rotation about each body's Z-axis.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        # Coincident origins
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        pos = [p_i[0] - p_j[0], p_i[1] - p_j[1], p_i[2] - p_j[2]]

        # Perpendicular Z-axes
        z_i = marker_z_axis(self.body_i, self.marker_i_quat)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)

        return pos + [dot3(z_i, z_j)]


# ============================================================================
# Phase 2: Mechanical element constraints
# ============================================================================


class GearConstraint(ConstraintBase):
    """Gear pair or belt — 1 DOF removed.

    Couples rotation angles: r_i * theta_i + r_j * theta_j = 0.
    For belts (same-direction rotation), r_j is passed as negative.

    Uses the Z-component of the relative quaternion as a proxy for
    rotation angle (linear for small angles, which is the regime
    where Newton operates).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_quat: tuple[float, float, float, float],
        radius_i: float,
        radius_j: float,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_quat = marker_i_quat
        self.marker_j_quat = marker_j_quat
        self.radius_i = radius_i
        self.radius_j = radius_j

    def residuals(self) -> List[Expr]:
        # Rotation angle proxy via relative quaternion Z-component
        # For body_i: q_rel_i = conj(q_marker_i) * q_body_i * q_marker_i
        # Simplified: use 2*qz of (conj(marker) * body * marker) as angle proxy
        qz_i = _rotation_z_component(self.body_i, self.marker_i_quat)
        qz_j = _rotation_z_component(self.body_j, self.marker_j_quat)

        # r_i * theta_i + r_j * theta_j = 0
        # Using qz as proportional to theta/2:
        # r_i * qz_i + r_j * qz_j = 0
        return [Const(self.radius_i) * qz_i + Const(self.radius_j) * qz_j]


class RackPinionConstraint(ConstraintBase):
    """Rack-and-pinion — 1 DOF removed.

    Couples rotation of body_i to translation of body_j along marker_j Z-axis.
    translation = pitch_radius * theta
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
        pitch_radius: float,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat
        self.pitch_radius = pitch_radius

    def residuals(self) -> List[Expr]:
        # Translation of j along its Z-axis
        p_i = self.body_i.world_point(*self.marker_i_pos)
        p_j = self.body_j.world_point(*self.marker_j_pos)
        z_j = marker_z_axis(self.body_j, self.marker_j_quat)
        d = sub3(p_j, p_i)
        translation = dot3(d, z_j)

        # Rotation angle of i about its Z-axis
        qz_i = _rotation_z_component(self.body_i, self.marker_i_quat)

        # translation - pitch_radius * angle = 0
        # angle ≈ 2 * qz, so: translation - pitch_radius * 2 * qz = 0
        return [translation - Const(2.0 * self.pitch_radius) * qz_i]


class CamConstraint(ConstraintBase):
    """Cam-follower constraint — future, stub."""

    def residuals(self) -> List[Expr]:
        return []


class SlotConstraint(ConstraintBase):
    """Slot constraint — future, stub."""

    def residuals(self) -> List[Expr]:
        return []


class DistanceCylSphConstraint(ConstraintBase):
    """Cylinder-sphere distance — stub.

    Semantics depend on geometry classification; placeholder for now.
    """

    def residuals(self) -> List[Expr]:
        return []


# -- rotation helpers for mechanical constraints ------------------------------


def _rotation_z_component(
    body: RigidBody,
    marker_quat: tuple[float, float, float, float],
) -> Expr:
    """Extract the Z-component of the relative quaternion about a marker axis.

    Returns the qz component of conj(q_marker) * q_body * q_marker,
    which is proportional to sin(theta/2) where theta is the rotation
    angle about the marker Z-axis.
    """
    mw, mx, my, mz = marker_quat
    # q_local = conj(marker) * q_body * marker
    # Step 1: temp = conj(marker) * q_body
    cmw, cmx, cmy, cmz = Const(mw), Const(-mx), Const(-my), Const(-mz)
    # temp = conj(marker) * q_body
    tw = cmw * body.qw - cmx * body.qx - cmy * body.qy - cmz * body.qz
    tx = cmw * body.qx + cmx * body.qw + cmy * body.qz - cmz * body.qy
    ty = cmw * body.qy - cmx * body.qz + cmy * body.qw + cmz * body.qx
    tz = cmw * body.qz + cmx * body.qy - cmy * body.qx + cmz * body.qw
    # q_local = temp * marker
    mmw, mmx, mmy, mmz = Const(mw), Const(mx), Const(my), Const(mz)
    # rz = tw * mmz + tx * mmy - ty * mmx + tz * mmw
    rz = tw * mmz + tx * mmy - ty * mmx + tz * mmw
    return rz


# -- quaternion multiplication helpers ----------------------------------------


def _quat_mul_const(
    aw: Expr,
    ax: Expr,
    ay: Expr,
    az: Expr,
    bw: float,
    bx: float,
    by: float,
    bz: float,
) -> tuple[Expr, Expr, Expr, Expr]:
    """Multiply Expr quaternion (a) by constant quaternion (b)."""
    cbw, cbx, cby, cbz = Const(bw), Const(bx), Const(by), Const(bz)
    return _quat_mul_expr(aw, ax, ay, az, cbw, cbx, cby, cbz)


def _quat_mul_expr(
    aw: Expr,
    ax: Expr,
    ay: Expr,
    az: Expr,
    bw: Expr,
    bx: Expr,
    by: Expr,
    bz: Expr,
) -> tuple[Expr, Expr, Expr, Expr]:
    """Hamilton product of two Expr quaternions."""
    rw = aw * bw - ax * bx - ay * by - az * bz
    rx = aw * bx + ax * bw + ay * bz - az * by
    ry = aw * by - ax * bz + ay * bw + az * bx
    rz = aw * bz + ax * by - ay * bx + az * bw
    return rw, rx, ry, rz