"""Geometric helper functions for constraint equations.

Provides Expr-level vector operations and marker axis extraction.
All functions work with Expr triples (tuples of 3 Expr nodes)
representing 3D vectors in world coordinates.

Marker convention (from Types.h): the marker's Z-axis defines the
constraint direction (hinge axis, face normal, line direction, etc.).
"""

from __future__ import annotations

from .entities import RigidBody
from .expr import Const, Expr
from .quat import quat_rotate

# Type alias for an Expr triple (3D vector)
Vec3 = tuple[Expr, Expr, Expr]


# -- Marker axis extraction ---------------------------------------------------


def _composed_quat(
    body: RigidBody,
    marker_quat: tuple[float, float, float, float],
) -> tuple[Expr, Expr, Expr, Expr]:
    """Compute q_total = q_body * q_marker as Expr quaternion.

    q_body comes from the body's Var params; q_marker is constant.
    """
    bw, bx, by, bz = body.qw, body.qx, body.qy, body.qz
    mw, mx, my, mz = (Const(v) for v in marker_quat)
    # Hamilton product: body * marker
    rw = bw * mw - bx * mx - by * my - bz * mz
    rx = bw * mx + bx * mw + by * mz - bz * my
    ry = bw * my - bx * mz + by * mw + bz * mx
    rz = bw * mz + bx * my - by * mx + bz * mw
    return rw, rx, ry, rz


def marker_z_axis(
    body: RigidBody,
    marker_quat: tuple[float, float, float, float],
) -> Vec3:
    """World-frame Z-axis of a marker on a body.

    Computes rotate(q_body * q_marker, [0, 0, 1]).
    """
    qw, qx, qy, qz = _composed_quat(body, marker_quat)
    return quat_rotate(qw, qx, qy, qz, Const(0.0), Const(0.0), Const(1.0))


def marker_x_axis(
    body: RigidBody,
    marker_quat: tuple[float, float, float, float],
) -> Vec3:
    """World-frame X-axis of a marker on a body.

    Computes rotate(q_body * q_marker, [1, 0, 0]).
    """
    qw, qx, qy, qz = _composed_quat(body, marker_quat)
    return quat_rotate(qw, qx, qy, qz, Const(1.0), Const(0.0), Const(0.0))


def marker_y_axis(
    body: RigidBody,
    marker_quat: tuple[float, float, float, float],
) -> Vec3:
    """World-frame Y-axis of a marker on a body.

    Computes rotate(q_body * q_marker, [0, 1, 0]).
    """
    qw, qx, qy, qz = _composed_quat(body, marker_quat)
    return quat_rotate(qw, qx, qy, qz, Const(0.0), Const(1.0), Const(0.0))


# -- Vector operations on Expr triples ----------------------------------------


def dot3(a: Vec3, b: Vec3) -> Expr:
    """Dot product of two Expr triples."""
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]


def cross3(a: Vec3, b: Vec3) -> Vec3:
    """Cross product of two Expr triples."""
    return (
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0],
    )


def sub3(a: Vec3, b: Vec3) -> Vec3:
    """Vector subtraction a - b."""
    return (a[0] - b[0], a[1] - b[1], a[2] - b[2])


# -- Geometric primitives -----------------------------------------------------


def point_plane_distance(
    point: Vec3,
    plane_origin: Vec3,
    normal: Vec3,
) -> Expr:
    """Signed distance from point to plane defined by origin + normal.

    Returns (point - plane_origin) . normal
    """
    d = sub3(point, plane_origin)
    return dot3(d, normal)


def point_line_perp_components(
    point: Vec3,
    line_origin: Vec3,
    line_dir: Vec3,
) -> tuple[Expr, Expr, Expr]:
    """Three perpendicular-distance components from point to line.

    The line passes through line_origin along line_dir.
    Returns all 3 components of (point - line_origin) x line_dir,
    which are zero when the point lies on the line.  Only 2 of 3 are
    independent, but using all 3 avoids a singularity when the line
    direction aligns with a coordinate axis.
    """
    d = sub3(point, line_origin)
    cx, cy, cz = cross3(d, line_dir)
    return cx, cy, cz