"""Monte Carlo statistical validation tests for AOA geolocation.

These tests verify that the solution covariance is statistically consistent
by checking that the normalized miss distance follows a chi-squared distribution.

Tests are skipped by default due to runtime. Run manually with:
    pytest test/geolocation/test_aoa_monte_carlo.py -v

Or run all tests including skipped:
    pytest test/geolocation/test_aoa_monte_carlo.py -v --run-monte-carlo
"""

from __future__ import annotations

import numpy as np
import pytest
from gri_geosim.geolocation import aoa_locator
from gri_utils.observables import get_aoa_quat
from scipy.spatial.transform import Rotation
from scipy.stats import chi2, kstest

# Skip all tests in this module by default
pytestmark = pytest.mark.skip(reason="Monte Carlo tests are slow; run manually")

EARTH_RADIUS = 6378137.0


def _create_good_geometry() -> tuple[np.ndarray, Rotation, np.ndarray]:
    """Create a test geometry with good angular separation.

    Returns collectors arranged in a triangle with ~120 degree separation,
    providing good observability in all directions.
    """
    # 3 collectors in a triangle, 20 km apart
    spread_km = 20.0
    angles = np.linspace(0, 2 * np.pi, 3, endpoint=False)

    collector_xyz = np.zeros((3, 3))
    for i, angle in enumerate(angles):
        collector_xyz[i] = [
            EARTH_RADIUS,
            spread_km * 1000 * np.cos(angle),
            spread_km * 1000 * np.sin(angle),
        ]

    # Identity rotations
    rotations = Rotation.identity(3)

    # Target 100 km away
    centroid = np.mean(collector_xyz, axis=0)
    direction = centroid / np.linalg.norm(centroid)
    target_xyz = centroid + direction * 100_000

    return collector_xyz, rotations, target_xyz


def test_normalized_miss_distribution() -> None:
    """Verify solution covariance is statistically consistent.

    Over many Monte Carlo trials, the normalized miss distance squared:
        d^2 = (x_est - x_true)^T @ P^{-1} @ (x_est - x_true)

    should follow a chi-squared(3) distribution (3 DOF for 3D position).
    """
    rng = np.random.default_rng(42)
    n_trials = 1000

    collector_xyz, rotations, target_true = _create_good_geometry()

    # AOA measurement noise: 1 mrad (1-sigma)
    aoa_sigma = 0.001
    aoa_cov = np.eye(2) * aoa_sigma**2
    aoa_covariances = np.tile(aoa_cov, (3, 1, 1))

    # Compute true AOA measurements (noiseless)
    true_aoa = get_aoa_quat(collector_xyz, target_true, rotations)

    normalized_miss_sq = []

    for _ in range(n_trials):
        # Add Gaussian noise to measurements
        noise = rng.standard_normal((3, 2)) * aoa_sigma
        noisy_aoa = true_aoa + noise

        # Solve
        result = aoa_locator(
            collector_xyz,
            rotations,
            noisy_aoa,
            aoa_covariances,
        )

        # Normalized miss distance squared
        delta = result.position - target_true
        cov_inv = np.linalg.inv(result.covariance)
        d_sq = delta @ cov_inv @ delta
        normalized_miss_sq.append(d_sq)

    normalized_miss_sq = np.array(normalized_miss_sq)

    # Chi-squared(3) statistics
    # Mean should be 3, variance should be 6
    mean_d_sq = np.mean(normalized_miss_sq)
    var_d_sq = np.var(normalized_miss_sq)

    # Allow statistical tolerance (with 1000 samples)
    assert abs(mean_d_sq - 3.0) < 0.3, f"Mean {mean_d_sq:.2f} not close to 3"
    assert abs(var_d_sq - 6.0) < 1.5, f"Variance {var_d_sq:.2f} not close to 6"

    # 95th percentile should be ~7.81 for chi2(3)
    p95 = np.percentile(normalized_miss_sq, 95)
    assert p95 < 10.0, f"95th percentile {p95:.2f} too high"

    # Kolmogorov-Smirnov test against chi2(3)
    _, p_value = kstest(normalized_miss_sq, chi2(3).cdf)
    assert p_value > 0.01, f"KS test p-value {p_value:.4f} too low"


def test_covariance_scaling_with_noise() -> None:
    """Verify covariance scales appropriately with measurement noise.

    Doubling measurement noise should roughly quadruple position variance.
    """
    rng = np.random.default_rng(123)
    n_trials = 500

    collector_xyz, rotations, target_true = _create_good_geometry()
    true_aoa = get_aoa_quat(collector_xyz, target_true, rotations)

    def run_trials(aoa_sigma: float) -> float:
        """Run trials and return mean squared position error."""
        aoa_cov = np.eye(2) * aoa_sigma**2
        aoa_covariances = np.tile(aoa_cov, (3, 1, 1))

        errors_sq = []
        for _ in range(n_trials):
            noise = rng.standard_normal((3, 2)) * aoa_sigma
            noisy_aoa = true_aoa + noise

            result = aoa_locator(
                collector_xyz,
                rotations,
                noisy_aoa,
                aoa_covariances,
            )
            error_sq = np.sum((result.position - target_true) ** 2)
            errors_sq.append(error_sq)

        return float(np.mean(errors_sq))

    # Compare 1 mrad vs 2 mrad noise
    mse_1mrad = run_trials(0.001)
    mse_2mrad = run_trials(0.002)

    # Expect ~4x increase (2^2 = 4)
    ratio = mse_2mrad / mse_1mrad
    assert 2.5 < ratio < 6.0, f"MSE ratio {ratio:.2f} not in expected range [2.5, 6.0]"


def test_poor_geometry_inflates_covariance() -> None:
    """Verify covariance inflates with poor collector geometry.

    Collectors nearly on a line should have much larger covariance compared
    to collectors in a well-distributed triangle, given similar baseline and range.
    """
    spread_km = 20.0  # Same baseline as good geometry
    target_range = 100_000  # Same range as good geometry

    # Poor geometry: collectors nearly collinear (small angular spread)
    # Clustered around Y-axis with minimal Z spread
    collector_xyz_poor = np.array(
        [
            [EARTH_RADIUS, -spread_km * 1000, 0.0],
            [EARTH_RADIUS, 0.0, 500.0],  # Small Z offset prevents singularity
            [EARTH_RADIUS, spread_km * 1000, 0.0],
        ],
    )
    rotations_poor = Rotation.identity(3)

    # Target 100km away
    centroid_poor = np.mean(collector_xyz_poor, axis=0)
    direction_poor = centroid_poor / np.linalg.norm(centroid_poor)
    target_poor = centroid_poor + direction_poor * target_range

    aoa_measurements_poor = get_aoa_quat(
        collector_xyz_poor,
        target_poor,
        rotations_poor,
    )
    aoa_cov = np.eye(2) * 0.001**2
    aoa_covariances_poor = np.tile(aoa_cov, (3, 1, 1))

    result_poor = aoa_locator(
        collector_xyz_poor,
        rotations_poor,
        aoa_measurements_poor,
        aoa_covariances_poor,
    )

    # Good geometry: collectors in a triangle (120-degree separation)
    collector_xyz_good, rotations_good, target_good = _create_good_geometry()
    aoa_measurements_good = get_aoa_quat(
        collector_xyz_good,
        target_good,
        rotations_good,
    )
    aoa_covariances_good = np.tile(aoa_cov, (3, 1, 1))

    result_good = aoa_locator(
        collector_xyz_good,
        rotations_good,
        aoa_measurements_good,
        aoa_covariances_good,
    )

    # Poor geometry should have larger trace (total variance)
    trace_poor = np.trace(result_poor.covariance)
    trace_good = np.trace(result_good.covariance)

    assert trace_poor > trace_good, (
        f"Poor geometry covariance ({trace_poor:.2e}) should exceed "
        f"good geometry ({trace_good:.2e})"
    )


def test_many_collectors_reduces_covariance() -> None:
    """Verify adding collectors reduces solution covariance."""
    aoa_sigma = 0.001
    aoa_cov = np.eye(2) * aoa_sigma**2

    traces = []
    for n_collectors in [2, 3, 5, 10]:
        # Create collectors in a circle
        angles = np.linspace(0, 2 * np.pi, n_collectors, endpoint=False)
        collector_xyz = np.zeros((n_collectors, 3))
        for i, angle in enumerate(angles):
            collector_xyz[i] = [
                EARTH_RADIUS,
                20000 * np.cos(angle),
                20000 * np.sin(angle),
            ]

        rotations = Rotation.identity(n_collectors)
        target_true = np.array([EARTH_RADIUS + 100000, 0.0, 0.0])

        aoa_measurements = get_aoa_quat(collector_xyz, target_true, rotations)
        aoa_covariances = np.tile(aoa_cov, (n_collectors, 1, 1))

        result = aoa_locator(
            collector_xyz,
            rotations,
            aoa_measurements,
            aoa_covariances,
        )
        traces.append(np.trace(result.covariance))

    # Each increase in collectors should reduce trace
    for i in range(len(traces) - 1):
        assert traces[i + 1] < traces[i], (
            f"Trace did not decrease: {traces[i]:.2e} -> {traces[i + 1]:.2e}"
        )
