from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import pytest
from scipy.stats import chi2

from gri_utils.ellipsoids import cep, cep_2sigma, cep_50, cep_90, cep_95, cep_sigma

if TYPE_CHECKING:
    from collections.abc import Callable

# Circular covariance: sigma = 5 in both axes
COV_CIRCULAR = np.array(((25.0, 0.0), (0.0, 25.0)))

# Elliptical covariance: sigma_x = 3, sigma_y = 10
COV_ELLIPTICAL = np.array(((9.0, 0.0), (0.0, 100.0)))

# 3x3 with circular horizontal component (sigma = 5), alt sigma = 8
COV_3D = np.array(((25.0, 0.0, 0.0), (0.0, 25.0, 0.0), (0.0, 0.0, 64.0)))


# -- circular case (Rayleigh): closed-form reference values --


@pytest.mark.parametrize(
    ("func", "expected"),
    [
        (cep_50, 5.0 * np.sqrt(np.log(4))),
        (cep_90, 5.0 * np.sqrt(-2.0 * np.log(0.10))),
        (cep_95, 5.0 * np.sqrt(-2.0 * np.log(0.05))),
    ],
)
def test_cep_circular(
    func: Callable[..., float],
    expected: float,
) -> None:
    assert func(COV_CIRCULAR) == pytest.approx(expected, rel=1e-10)


def test_cep_sigma_circular() -> None:
    # 1-sigma for 2D Rayleigh: conf = 1 - exp(-0.5) ~= 0.3935
    expected = 5.0 * np.sqrt(-2.0 * np.log(np.exp(-0.5)))
    assert cep_sigma(COV_CIRCULAR) == pytest.approx(expected, rel=1e-8)


def test_cep_2sigma_circular() -> None:
    conf = float(chi2.cdf(4.0, df=2))
    expected = 5.0 * np.sqrt(-2.0 * np.log(1.0 - conf))
    assert cep_2sigma(COV_CIRCULAR) == pytest.approx(expected, rel=1e-8)


# -- elliptical case: verify against Monte Carlo --


def test_cep_elliptical_monte_carlo() -> None:
    rng = np.random.default_rng(42)
    samples = rng.multivariate_normal([0, 0], COV_ELLIPTICAL, size=500_000)
    radii = np.linalg.norm(samples, axis=1)
    mc_cep50 = float(np.median(radii))
    mc_cep95 = float(np.percentile(radii, 95))
    assert cep_50(COV_ELLIPTICAL) == pytest.approx(mc_cep50, rel=0.01)
    assert cep_95(COV_ELLIPTICAL) == pytest.approx(mc_cep95, rel=0.01)


# -- 3x3 input: projects to 2x2 --


def test_cep_3d_matches_2d() -> None:
    assert cep_50(COV_3D) == pytest.approx(cep_50(COV_CIRCULAR), rel=1e-10)


# -- monotonicity and scaling --


def test_cep_monotonicity() -> None:
    r50 = cep(COV_ELLIPTICAL, 0.50)
    r90 = cep(COV_ELLIPTICAL, 0.90)
    r95 = cep(COV_ELLIPTICAL, 0.95)
    assert r50 < r90 < r95


def test_cep_scales_with_sigma() -> None:
    cov_scaled = 4.0 * COV_ELLIPTICAL  # sigma doubled -> cov x4
    assert cep_50(cov_scaled) == pytest.approx(2.0 * cep_50(COV_ELLIPTICAL), rel=1e-6)


# -- error handling --


def test_cep_invalid_conf() -> None:
    with pytest.raises(ValueError, match="conf"):
        cep(COV_CIRCULAR, 0.0)
    with pytest.raises(ValueError, match="conf"):
        cep(COV_CIRCULAR, 1.0)


def test_cep_invalid_shape() -> None:
    with pytest.raises(ValueError, match="shape"):
        cep(np.eye(4))


def test_cep_accepts_sequences() -> None:
    seq = ((25.0, 0.0), (0.0, 25.0))
    assert cep_50(seq) == pytest.approx(cep_50(COV_CIRCULAR), rel=1e-10)


# -- extreme sigma ratios (i0e stability) --


@pytest.mark.parametrize(
    ("sig1", "sig2"),
    [(1.0, 100.0), (1.0, 1000.0)],
    ids=["1:100", "1:1000"],
)
def test_cep_extreme_ratio_monte_carlo(sig1: float, sig2: float) -> None:
    cov = np.diag([sig1**2, sig2**2])
    rng = np.random.default_rng(42)
    samples = rng.multivariate_normal([0, 0], cov, size=2_000_000)
    radii = np.linalg.norm(samples, axis=1)
    mc_cep50 = float(np.median(radii))
    mc_cep95 = float(np.percentile(radii, 95))
    assert cep_50(cov) == pytest.approx(mc_cep50, rel=0.01)
    assert cep_95(cov) == pytest.approx(mc_cep95, rel=0.01)
