"""Quaternion math expressed as Expr trees.

All functions take Expr nodes (typically Var nodes from a RigidBody's
parameter set) and return Expr trees.  Quaternion rotation is polynomial
in the quaternion components — no transcendentals needed.

Convention: quaternion (qw, qx, qy, qz) matches KCSolve's (w, x, y, z).
"""

from __future__ import annotations

from .expr import Const, Expr


def quat_rotate(
    qw: Expr,
    qx: Expr,
    qy: Expr,
    qz: Expr,
    px: Expr,
    py: Expr,
    pz: Expr,
) -> tuple[Expr, Expr, Expr]:
    """Rotate point (px, py, pz) by unit quaternion (qw, qx, qy, qz).

    Uses the standard formula:  p' = q * p * q^{-1}
    expanded into component form (polynomial in q):

        rx = (1 - 2(qy^2 + qz^2)) * px + 2(qx*qy - qw*qz) * py + 2(qx*qz + qw*qy) * pz
        ry = 2(qx*qy + qw*qz) * px + (1 - 2(qx^2 + qz^2)) * py + 2(qy*qz - qw*qx) * pz
        rz = 2(qx*qz - qw*qy) * px + 2(qy*qz + qw*qx) * py + (1 - 2(qx^2 + qy^2)) * pz
    """
    two = Const(2.0)

    rx = (
        (Const(1.0) - two * (qy * qy + qz * qz)) * px
        + two * (qx * qy - qw * qz) * py
        + two * (qx * qz + qw * qy) * pz
    )
    ry = (
        two * (qx * qy + qw * qz) * px
        + (Const(1.0) - two * (qx * qx + qz * qz)) * py
        + two * (qy * qz - qw * qx) * pz
    )
    rz = (
        two * (qx * qz - qw * qy) * px
        + two * (qy * qz + qw * qx) * py
        + (Const(1.0) - two * (qx * qx + qy * qy)) * pz
    )
    return rx, ry, rz


def world_point(
    tx: Expr,
    ty: Expr,
    tz: Expr,
    qw: Expr,
    qx: Expr,
    qy: Expr,
    qz: Expr,
    lx: float,
    ly: float,
    lz: float,
) -> tuple[Expr, Expr, Expr]:
    """Transform local point (lx, ly, lz) to world coordinates.

    world = rotation(q, local) + translation
    """
    clx, cly, clz = Const(lx), Const(ly), Const(lz)
    rx, ry, rz = quat_rotate(qw, qx, qy, qz, clx, cly, clz)
    return tx + rx, ty + ry, tz + rz