festifun/backend/festival4d/geometry.py

"""Coordinate conversions, quaternion math, ray casting, triangulation.

The single most important thing in this file is :func:`colmap_to_threejs` — the
COLMAP world->camera pose to Three.js camera conversion (spec M5). It is a **FROZEN
CONTRACT**: it ships in foundation with a passing unit test (``test_geometry.py``), and
the JS mirror lives in ``frontend/src/lib/pose.js`` with the *same* embedded test
vectors. Lanes B and C **consume** it; they never reimplement or modify it.

The remaining functions (:func:`slerp_pose`, :func:`ray_from_pixel`,
:func:`triangulate_rays`, :func:`nearest_point_on_ray`) are stubs with frozen
signatures for lane B (spec M2 + M8). Lane B fills the bodies and adds their tests.

Conventions
-----------
- Quaternions are COLMAP order **[w, x, y, z]** (scalar first), unit norm.
- A pose ``(q, t)`` is world->camera: ``x_cam = R(q) @ x_world + t``.
- COLMAP camera axes: +x right, +y down, +z forward (into the scene).
- Three.js cameras look down -z with +y up; hence the ``diag(1, -1, -1)`` flip.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import ArrayLike, NDArray

# Camera-axis flip that takes COLMAP camera-local axes (x right, y down, z forward)
# to Three.js camera-local axes (x right, y up, z backward). Part of the frozen contract.
_FLIP_YZ = np.diag([1.0, -1.0, -1.0])


def quat_to_mat(q: ArrayLike) -> NDArray[np.float64]:
    """Convert a unit quaternion ``[w, x, y, z]`` to a 3x3 rotation matrix.

    Hamilton convention, right-handed, active rotation: the returned ``R`` is COLMAP's
    world->camera matrix when ``q`` is a COLMAP pose quaternion. FROZEN (used by the
    conversion contract).
    """
    w, x, y, z = (float(v) for v in np.asarray(q, dtype=np.float64).reshape(4))
    n = w * w + x * x + y * y + z * z
    if n < 1e-12:
        raise ValueError("quaternion has near-zero norm")
    s = 2.0 / n
    wx, wy, wz = s * w * x, s * w * y, s * w * z
    xx, xy, xz = s * x * x, s * x * y, s * x * z
    yy, yz, zz = s * y * y, s * y * z, s * z * z
    return np.array(
        [
            [1.0 - (yy + zz), xy - wz, xz + wy],
            [xy + wz, 1.0 - (xx + zz), yz - wx],
            [xz - wy, yz + wx, 1.0 - (xx + yy)],
        ],
        dtype=np.float64,
    )


def mat_to_quat(R: ArrayLike) -> NDArray[np.float64]:
    """Convert a 3x3 rotation matrix to a unit quaternion ``[w, x, y, z]`` (w >= 0).

    Inverse of :func:`quat_to_mat`. Used by tests and the pose export path.
    """
    m = np.asarray(R, dtype=np.float64).reshape(3, 3)
    trace = m[0, 0] + m[1, 1] + m[2, 2]
    if trace > 0.0:
        s = np.sqrt(trace + 1.0) * 2.0
        w = 0.25 * s
        x = (m[2, 1] - m[1, 2]) / s
        y = (m[0, 2] - m[2, 0]) / s
        z = (m[1, 0] - m[0, 1]) / s
    elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
        s = np.sqrt(1.0 + m[0, 0] - m[1, 1] - m[2, 2]) * 2.0
        w = (m[2, 1] - m[1, 2]) / s
        x = 0.25 * s
        y = (m[0, 1] + m[1, 0]) / s
        z = (m[0, 2] + m[2, 0]) / s
    elif m[1, 1] > m[2, 2]:
        s = np.sqrt(1.0 + m[1, 1] - m[0, 0] - m[2, 2]) * 2.0
        w = (m[0, 2] - m[2, 0]) / s
        x = (m[0, 1] + m[1, 0]) / s
        y = 0.25 * s
        z = (m[1, 2] + m[2, 1]) / s
    else:
        s = np.sqrt(1.0 + m[2, 2] - m[0, 0] - m[1, 1]) * 2.0
        w = (m[1, 0] - m[0, 1]) / s
        x = (m[0, 2] + m[2, 0]) / s
        y = (m[1, 2] + m[2, 1]) / s
        z = 0.25 * s
    q = np.array([w, x, y, z], dtype=np.float64)
    q /= np.linalg.norm(q)
    if q[0] < 0:  # canonical sign: non-negative scalar part
        q = -q
    return q


def colmap_to_threejs(q: ArrayLike, t: ArrayLike) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Convert a COLMAP world->camera pose to a Three.js camera pose (spec M5).

    **FROZEN CONTRACT.** Mirrored in ``frontend/src/lib/pose.js``; do not change the math
    without a change request and a synchronized update to both sides + their test vectors.

    Parameters
    ----------
    q : array_like, shape (4,)
        COLMAP world->camera quaternion ``[w, x, y, z]``.
    t : array_like, shape (3,)
        COLMAP world->camera translation ``[tx, ty, tz]``.

    Returns
    -------
    position : ndarray, shape (3,)
        Camera center in world coordinates, ``C = -R^T t``. Assign to ``camera.position``.
    rotation_matrix : ndarray, shape (3, 3)
        Three.js camera world rotation ``R_three = R^T @ diag(1, -1, -1)``. Assign via
        ``camera.setRotationFromMatrix(...)`` (a proper rotation, det = +1).
    """
    R = quat_to_mat(q)                      # world -> cam
    t_vec = np.asarray(t, dtype=np.float64).reshape(3)
    R_c2w = R.T                             # cam -> world
    position = -R_c2w @ t_vec              # camera center in world coords
    rotation_matrix = R_c2w @ _FLIP_YZ     # flip camera-local y,z for Three.js
    return position, rotation_matrix


# ---------------------------------------------------------------------------
# Lane B (spec M2 + M8): pose interpolation, pixel ray casting, two-view triangulation,
# and the single-view nearest-point fallback. Signatures FROZEN; bodies implemented below.
# ---------------------------------------------------------------------------
def slerp_pose(
    q0: ArrayLike,
    t0: ArrayLike,
    q1: ArrayLike,
    t1: ArrayLike,
    alpha: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Interpolate between two world->camera poses (spec M2 pose interpolation).

    Spherically interpolate the rotation (slerp on the shorter arc, handling the
    double-cover sign) and linearly interpolate the translation, at fraction
    ``alpha in [0, 1]`` from pose 0 to pose 1.

    Parameters
    ----------
    q0, q1 : array_like, shape (4,)
        COLMAP quaternions ``[w, x, y, z]`` at the endpoints.
    t0, t1 : array_like, shape (3,)
        COLMAP translations at the endpoints.
    alpha : float
        Interpolation fraction; 0 returns pose 0, 1 returns pose 1.

    Returns
    -------
    (q, t) : the interpolated quaternion ``[w, x, y, z]`` and translation ``[x, y, z]``.
    """
    a = float(alpha)
    q0 = np.asarray(q0, dtype=np.float64).reshape(4)
    q1 = np.asarray(q1, dtype=np.float64).reshape(4)
    n0 = np.linalg.norm(q0)
    n1 = np.linalg.norm(q1)
    if n0 < 1e-12 or n1 < 1e-12:
        raise ValueError("slerp endpoint quaternion has near-zero norm")
    q0 = q0 / n0
    q1 = q1 / n1

    # Double cover: pick the sign of q1 that lies on the same hemisphere as q0, so slerp
    # takes the shorter arc (a rotation and its negation are the same orientation).
    dot = float(np.dot(q0, q1))
    if dot < 0.0:
        q1 = -q1
        dot = -dot
    dot = min(1.0, max(-1.0, dot))

    if dot > 0.9995:
        # Endpoints almost coincide: nlerp is numerically safe and visually identical.
        q_interp = q0 + a * (q1 - q0)
    else:
        theta_0 = np.arccos(dot)
        sin_0 = np.sin(theta_0)
        s0 = np.sin((1.0 - a) * theta_0) / sin_0
        s1 = np.sin(a * theta_0) / sin_0
        q_interp = s0 * q0 + s1 * q1
    q_interp = q_interp / np.linalg.norm(q_interp)

    t0 = np.asarray(t0, dtype=np.float64).reshape(3)
    t1 = np.asarray(t1, dtype=np.float64).reshape(3)
    t_interp = (1.0 - a) * t0 + a * t1
    return q_interp, t_interp


def ray_from_pixel(
    q: ArrayLike,
    t: ArrayLike,
    fx: float,
    fy: float,
    cx: float,
    cy: float,
    px: float,
    py: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Unproject a pixel into a world-space ray (spec M8 annotation resolution).

    Given the world->camera pose ``(q, t)`` and pinhole intrinsics, build the ray that
    passes through image pixel ``(px, py)``.

    Returns
    -------
    origin : ndarray, shape (3,)
        Ray origin = camera center in world coords.
    direction : ndarray, shape (3,)
        Unit ray direction in world coords, pointing into the scene.
    """
    R = quat_to_mat(q)                          # world -> cam
    t_vec = np.asarray(t, dtype=np.float64).reshape(3)
    origin = -R.T @ t_vec                        # camera center in world coords

    # Pinhole back-projection. A world point X projects with x_cam = R X + t and
    # px = fx * x_cam.x / x_cam.z + cx, py = fy * x_cam.y / x_cam.z + cy
    # (COLMAP camera axes: +x right, +y down, +z forward). So the camera-space direction
    # through pixel (px, py) is [(px-cx)/fx, (py-cy)/fy, 1], pointing forward into the scene.
    d_cam = np.array([(px - cx) / fx, (py - cy) / fy, 1.0], dtype=np.float64)
    d_world = R.T @ d_cam                         # rotate direction cam -> world
    norm = np.linalg.norm(d_world)
    if norm < 1e-12:
        raise ValueError("degenerate ray direction")
    return origin, d_world / norm


def triangulate_rays(
    origin_a: ArrayLike,
    dir_a: ArrayLike,
    origin_b: ArrayLike,
    dir_b: ArrayLike,
) -> tuple[NDArray[np.float64], float]:
    """Closest point between two world-space rays (spec M8 two-view triangulation).

    Returns
    -------
    point : ndarray, shape (3,)
        Midpoint of the shortest segment connecting the two rays.
    gap : float
        Length of that shortest segment (the mutual-nearest-approach distance). Callers
        reject the triangulation when the rays are near-parallel or ``gap`` exceeds the
        spec threshold (0.5 scene units).
    """
    oa = np.asarray(origin_a, dtype=np.float64).reshape(3)
    ob = np.asarray(origin_b, dtype=np.float64).reshape(3)
    da = np.asarray(dir_a, dtype=np.float64).reshape(3)
    db = np.asarray(dir_b, dtype=np.float64).reshape(3)
    na, nb = np.linalg.norm(da), np.linalg.norm(db)
    if na < 1e-12 or nb < 1e-12:
        raise ValueError("triangulate_rays: zero-length direction")
    da = da / na
    db = db / nb

    # Shortest segment between two lines P(s)=oa+s*da, Q(u)=ob+u*db. Minimize |P-Q|^2.
    # With unit directions: b = da.db, denom = 1 - b^2 (0 when parallel).
    w0 = oa - ob
    b = float(np.dot(da, db))
    d = float(np.dot(da, w0))
    e = float(np.dot(db, w0))
    denom = 1.0 - b * b
    if denom < 1e-9:
        # Near-parallel: no unique closest pair. Anchor on oa, take the closest point on
        # line b to it; gap is the line-to-line perpendicular distance. Callers reject
        # near-parallel rays up front, so this branch just stays numerically safe.
        s = 0.0
        u = e
    else:
        s = (b * e - d) / denom
        u = (e - b * d) / denom
    pa = oa + s * da
    pb = ob + u * db
    point = 0.5 * (pa + pb)
    gap = float(np.linalg.norm(pa - pb))
    return point, gap


def nearest_point_on_ray(
    origin: ArrayLike,
    direction: ArrayLike,
    points: ArrayLike,
    radius: float = 0.3,
) -> NDArray[np.float64] | None:
    """Nearest point-cloud point to a ray, within a cylinder (spec M8 single-view fallback).

    Among ``points`` (shape ``(N, 3)``) find the one whose perpendicular distance to the
    ray is smallest, considering only points within ``radius`` of the ray and in front of
    the origin.

    Returns
    -------
    point : ndarray shape (3,) or None
        The selected point-cloud point, or ``None`` if none lie within ``radius``.
    """
    o = np.asarray(origin, dtype=np.float64).reshape(3)
    d = np.asarray(direction, dtype=np.float64).reshape(3)
    nd = np.linalg.norm(d)
    if nd < 1e-12:
        raise ValueError("nearest_point_on_ray: zero-length direction")
    d = d / nd

    pts = np.asarray(points, dtype=np.float64).reshape(-1, 3)
    if len(pts) == 0:
        return None

    v = pts - o                          # origin -> each point
    proj = v @ d                         # signed distance along the ray
    perp = v - np.outer(proj, d)         # component perpendicular to the ray
    perp_dist = np.linalg.norm(perp, axis=1)

    # In front of the origin and inside the cylinder of the given radius.
    mask = (proj > 0.0) & (perp_dist <= radius)
    if not np.any(mask):
        return None
    idx_in = np.where(mask)[0]
    best = idx_in[np.argmin(perp_dist[idx_in])]
    return pts[best].copy()