"""Tests for geometry.py.

The ``colmap_to_threejs`` tests below are a **FROZEN CONTRACT** (foundation). The embedded
``POSE_TEST_VECTORS`` are duplicated verbatim in ``frontend/src/lib/pose.js`` so both sides
provably agree — if you touch the conversion math, update both files and both vector sets,
via a change request. Lane B owns this file but must not modify these three cases; it adds
tests for slerp_pose / ray_from_pixel / triangulate_rays / nearest_point_on_ray below them.
"""

from __future__ import annotations

import numpy as np
import pytest
from scipy.spatial.transform import Rotation

from festival4d.geometry import (
    colmap_to_threejs,
    mat_to_quat,
    quat_to_mat,
)

# Sqrt(2)/2, spelled out so it matches the JS constant character-for-character.
SQRT1_2 = 0.7071067811865476

# --- FROZEN test vectors (must equal POSE_TEST_VECTORS in frontend/src/lib/pose.js) ----
# Each: COLMAP world->camera (q=[w,x,y,z], t) -> Three.js (position, rotation 3x3 row-major).
POSE_TEST_VECTORS = [
    {
        "name": "identity",
        "q": [1.0, 0.0, 0.0, 0.0],
        "t": [0.0, 0.0, -10.0],
        "position": [0.0, 0.0, 10.0],
        "rotation": [[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]],
    },
    {
        "name": "yaw90",  # +90 deg about world Y
        "q": [SQRT1_2, 0.0, SQRT1_2, 0.0],
        "t": [0.0, 0.0, 10.0],
        "position": [10.0, 0.0, 0.0],
        "rotation": [[0.0, 0.0, 1.0], [0.0, -1.0, 0.0], [1.0, 0.0, 0.0]],
    },
    {
        "name": "lookat",  # camera at +Z looking at origin -> Three.js rotation is identity
        "q": [0.0, 1.0, 0.0, 0.0],
        "t": [0.0, 0.0, 8.0],
        "position": [0.0, 0.0, 8.0],
        "rotation": [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
    },
]


@pytest.mark.parametrize("vec", POSE_TEST_VECTORS, ids=lambda v: v["name"])
def test_colmap_to_threejs_known_vectors(vec):
    position, rotation = colmap_to_threejs(vec["q"], vec["t"])
    np.testing.assert_allclose(position, vec["position"], atol=1e-9)
    np.testing.assert_allclose(rotation, vec["rotation"], atol=1e-9)


@pytest.mark.parametrize("vec", POSE_TEST_VECTORS, ids=lambda v: v["name"])
def test_threejs_rotation_is_proper(vec):
    """R_three must be a proper rotation (orthonormal, det = +1)."""
    _, R = colmap_to_threejs(vec["q"], vec["t"])
    np.testing.assert_allclose(R @ R.T, np.eye(3), atol=1e-9)
    assert abs(np.linalg.det(R) - 1.0) < 1e-9


def _colmap_pose(R_w2c: np.ndarray, center: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Build a COLMAP (q, t) from a world->camera rotation and a world camera center."""
    t = -R_w2c @ center
    q = mat_to_quat(R_w2c)
    return q, t


def test_colmap_to_threejs_roundtrip_random_poses():
    """Forward COLMAP->Three.js then invert it; recover the original (q, t)."""
    rng = np.random.default_rng(20260716)
    flip = np.diag([1.0, -1.0, -1.0])
    for _ in range(50):
        R_w2c = Rotation.random(random_state=rng).as_matrix()
        center = rng.uniform(-10, 10, size=3)
        q, t = _colmap_pose(R_w2c, center)

        position, R_three = colmap_to_threejs(q, t)

        # position is the camera center
        np.testing.assert_allclose(position, center, atol=1e-9)

        # invert: R_c2w = R_three @ flip  ->  R_w2c = R_c2w^T  ->  t = -R_w2c @ C
        R_c2w = R_three @ flip
        R_back = R_c2w.T
        t_back = -R_back @ position
        q_back = mat_to_quat(R_back)

        np.testing.assert_allclose(t_back, t, atol=1e-9)
        # quaternions are double-cover; compare canonical (w>=0) forms
        q_canon = q if q[0] >= 0 else -q
        np.testing.assert_allclose(q_back, q_canon, atol=1e-9)


def test_quat_to_mat_matches_scipy_oracle():
    """quat_to_mat agrees with an independent implementation (scipy)."""
    rng = np.random.default_rng(11)
    for _ in range(50):
        # scipy quaternions are scalar-LAST [x,y,z,w]; ours are scalar-FIRST [w,x,y,z]
        q_xyzw = rng.normal(size=4)
        q_xyzw /= np.linalg.norm(q_xyzw)
        q_wxyz = np.array([q_xyzw[3], q_xyzw[0], q_xyzw[1], q_xyzw[2]])
        np.testing.assert_allclose(
            quat_to_mat(q_wxyz), Rotation.from_quat(q_xyzw).as_matrix(), atol=1e-9
        )