"""AZ-920 — `recover_scale` pure-function tests.

Pins the GSD scale-recovery contract that the KLT/RANSAC integration
in ``KltRansacStrategy.process_frame`` relies on. The integration
test is in ``test_az920_klt_ransac_scale_integration.py``; this file
covers the math in isolation.
"""

from __future__ import annotations

import math

import numpy as np
import pytest

from gps_denied_onboard.components.c1_vio._gsd_scale import (
    INFINITE_SCALE_UNCERTAINTY_M,
    RELATIVE_DISPLACEMENT_NOISE,
    ScaleRecoveryResult,
    recover_scale,
)

# ----------------------------------------------------------------------
# Helpers


def _intrinsics(fx: float = 1680.4469) -> np.ndarray:
    """Derkachi-grade nadir intrinsics by default (fx=fy, cx/cy at centre)."""
    return np.array(
        [
            [fx, 0.0, 960.0],
            [0.0, fx, 540.0],
            [0.0, 0.0, 1.0],
        ],
        dtype=np.float64,
    )


def _uniform_flow(n: int, dx_px: float, dy_px: float = 0.0) -> tuple[np.ndarray, np.ndarray]:
    """Build Nx2 prev / curr point clouds with a uniform pixel-flow shift."""
    rng = np.random.default_rng(seed=2026)
    base = rng.uniform(100.0, 1800.0, size=(n, 2))
    prev = base.copy()
    curr = base + np.array([dx_px, dy_px], dtype=np.float64)
    return prev, curr


# ----------------------------------------------------------------------
# Happy path


def test_nadir_identity_rotation_with_known_gsd_yields_expected_scale() -> None:
    # Arrange — at AGL=100 m, fx=1680 px → GSD ≈ 0.0595 m/px. A 20 px
    # mean flow corresponds to ≈1.19 m of horizontal motion. t_unit
    # along +X means t_horizontal = 1; scale_factor must equal 1.19.
    t_unit = np.array([1.0, 0.0, 0.0], dtype=np.float64)
    pts_prev, pts_curr = _uniform_flow(n=64, dx_px=20.0)

    # Act
    result = recover_scale(
        t_unit=t_unit,
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(),
        agl_m=100.0,
    )

    # Assert
    assert isinstance(result, ScaleRecoveryResult)
    expected_disp_m = 20.0 * 100.0 / 1680.4469
    assert result.scale_factor == pytest.approx(expected_disp_m, rel=1e-9)
    assert result.scale_uncertainty_m == pytest.approx(
        expected_disp_m * RELATIVE_DISPLACEMENT_NOISE, rel=1e-9
    )
    assert result.is_recoverable()
    # AZ-921 — happy-path categorisation is "metric".
    assert result.quality == "metric"


def test_diagonal_flow_uses_l2_norm_for_mean_magnitude() -> None:
    # Arrange — (dx=3, dy=4) flow has L2 magnitude 5 px per feature.
    t_unit = np.array([1.0, 0.0, 0.0], dtype=np.float64)
    pts_prev, pts_curr = _uniform_flow(n=50, dx_px=3.0, dy_px=4.0)

    # Act
    result = recover_scale(
        t_unit=t_unit,
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(),
        agl_m=100.0,
    )

    # Assert
    expected_disp_m = 5.0 * 100.0 / 1680.4469
    assert result.scale_factor == pytest.approx(expected_disp_m, rel=1e-9)


def test_off_axis_t_unit_uses_horizontal_component_magnitude() -> None:
    # Arrange — t_unit splits 0.6/0.8 between camera-X and camera-Y;
    # the horizontal magnitude is sqrt(0.36 + 0.64) = 1.0. The scale
    # therefore stays the displacement / 1.0.
    t_unit = np.array([0.6, 0.8, 0.0], dtype=np.float64)
    pts_prev, pts_curr = _uniform_flow(n=40, dx_px=10.0)

    # Act
    result = recover_scale(
        t_unit=t_unit,
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(),
        agl_m=120.0,
    )

    # Assert
    expected_disp_m = 10.0 * 120.0 / 1680.4469
    assert result.scale_factor == pytest.approx(expected_disp_m, rel=1e-9)


# ----------------------------------------------------------------------
# Degenerate / fallback paths — each must produce inf uncertainty.


def test_agl_none_returns_infinite_uncertainty() -> None:
    pts_prev, pts_curr = _uniform_flow(n=32, dx_px=20.0)

    result = recover_scale(
        t_unit=np.array([1.0, 0.0, 0.0]),
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(),
        agl_m=None,
    )

    assert result.scale_factor == 0.0
    assert result.scale_uncertainty_m == INFINITE_SCALE_UNCERTAINTY_M
    assert not result.is_recoverable()
    # AZ-921 — AGL-unknown is "unknown", not "direction_only".
    assert result.quality == "unknown"


def test_zero_or_negative_agl_returns_infinite_uncertainty() -> None:
    pts_prev, pts_curr = _uniform_flow(n=32, dx_px=20.0)

    for bad_agl in (0.0, -1.0):
        result = recover_scale(
            t_unit=np.array([1.0, 0.0, 0.0]),
            pts_prev=pts_prev,
            pts_curr=pts_curr,
            intrinsics_3x3=_intrinsics(),
            agl_m=bad_agl,
        )

        assert math.isinf(result.scale_uncertainty_m), f"agl={bad_agl}"


def test_near_pure_vertical_t_unit_returns_infinite_uncertainty() -> None:
    # Arrange — t along the optical axis (camera-Z) only; the GSD
    # method cannot tie horizontal pixel flow to vertical motion.
    pts_prev, pts_curr = _uniform_flow(n=40, dx_px=5.0)

    result = recover_scale(
        t_unit=np.array([0.01, 0.01, 0.9998]),
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(),
        agl_m=100.0,
    )

    assert not result.is_recoverable()
    # AZ-921 — near-vertical motion routes to direction_only: the
    # direction of t_unit is still meaningful, only the magnitude is
    # not derivable from horizontal pixel flow.
    assert result.quality == "direction_only"


def test_zero_inlier_flow_returns_infinite_uncertainty() -> None:
    # Arrange — stationary inliers (dx=dy=0). Cannot disambiguate the
    # scale of a zero motion.
    pts_prev, pts_curr = _uniform_flow(n=40, dx_px=0.0, dy_px=0.0)

    result = recover_scale(
        t_unit=np.array([1.0, 0.0, 0.0]),
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(),
        agl_m=100.0,
    )

    assert not result.is_recoverable()
    # AZ-921 — stale features routes to "unknown", since neither the
    # magnitude nor the direction of motion is trustworthy.
    assert result.quality == "unknown"


def test_empty_inlier_set_returns_infinite_uncertainty() -> None:
    empty = np.zeros((0, 2), dtype=np.float64)

    result = recover_scale(
        t_unit=np.array([1.0, 0.0, 0.0]),
        pts_prev=empty,
        pts_curr=empty,
        intrinsics_3x3=_intrinsics(),
        agl_m=100.0,
    )

    assert not result.is_recoverable()


def test_zero_focal_length_returns_infinite_uncertainty() -> None:
    pts_prev, pts_curr = _uniform_flow(n=32, dx_px=20.0)

    result = recover_scale(
        t_unit=np.array([1.0, 0.0, 0.0]),
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(fx=0.0),
        agl_m=100.0,
    )

    assert not result.is_recoverable()


# ----------------------------------------------------------------------
# Numerical stability with a Derkachi-grade inlier set


def test_derkachi_grade_inlier_count_yields_stable_scale() -> None:
    # Arrange — 120 inliers (representative of the Derkachi cruise
    # observation) with a small per-feature noise around a uniform
    # 30 px optical flow at 150 m AGL.
    rng = np.random.default_rng(seed=271828)
    base = rng.uniform(100.0, 1800.0, size=(120, 2))
    flow_mean_px = 30.0
    noise = rng.normal(0.0, 0.5, size=(120, 2))
    pts_prev = base
    pts_curr = base + np.array([flow_mean_px, 0.0]) + noise
    t_unit = np.array([1.0, 0.0, 0.0])
    agl_m = 150.0

    # Act
    result = recover_scale(
        t_unit=t_unit,
        pts_prev=pts_prev,
        pts_curr=pts_curr,
        intrinsics_3x3=_intrinsics(),
        agl_m=agl_m,
    )

    # Assert — observed scale must be within 5 % of the analytic value
    # at this noise level (sigma = 0.5 px over 120 inliers).
    expected = flow_mean_px * agl_m / 1680.4469
    assert result.scale_factor == pytest.approx(expected, rel=5e-2)
    assert result.is_recoverable()