from __future__ import annotations

from typing import TYPE_CHECKING

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

from gri_utils.ellipsoids import sep, sep_2sigma, sep_50, sep_90, sep_95, sep_sigma

if TYPE_CHECKING:
    from collections.abc import Callable

# Spherical covariance: sigma = 5 in all axes
COV_SPHERICAL = np.diag([25.0, 25.0, 25.0])

# Ellipsoidal covariance: sigma = 3, 5, 10
COV_ELLIPSOIDAL = np.diag([9.0, 25.0, 100.0])


# -- spherical case (chi distribution k=3): closed-form reference --


@pytest.mark.parametrize(
    ("func", "conf"),
    [(sep_50, 0.50), (sep_90, 0.90), (sep_95, 0.95)],
)
def test_sep_spherical(func: Callable[..., float], conf: float) -> None:
    expected = 5.0 * float(chi.ppf(conf, df=3))
    assert func(COV_SPHERICAL) == pytest.approx(expected, rel=1e-10)


def test_sep_sigma_spherical() -> None:
    conf = float(chi2.cdf(1.0, df=3))
    expected = 5.0 * float(chi.ppf(conf, df=3))
    assert sep_sigma(COV_SPHERICAL) == pytest.approx(expected, rel=1e-8)


def test_sep_2sigma_spherical() -> None:
    conf = float(chi2.cdf(4.0, df=3))
    expected = 5.0 * float(chi.ppf(conf, df=3))
    assert sep_2sigma(COV_SPHERICAL) == pytest.approx(expected, rel=1e-8)


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


def test_sep_ellipsoidal_monte_carlo() -> None:
    rng = np.random.default_rng(42)
    samples = rng.multivariate_normal([0, 0, 0], COV_ELLIPSOIDAL, size=500_000)
    radii = np.linalg.norm(samples, axis=1)
    mc_sep50 = float(np.median(radii))
    mc_sep95 = float(np.percentile(radii, 95))
    assert sep_50(COV_ELLIPSOIDAL) == pytest.approx(mc_sep50, rel=0.01)
    assert sep_95(COV_ELLIPSOIDAL) == pytest.approx(mc_sep95, rel=0.01)


# -- monotonicity and scaling --


def test_sep_monotonicity() -> None:
    r50 = sep(COV_ELLIPSOIDAL, 0.50)
    r90 = sep(COV_ELLIPSOIDAL, 0.90)
    r95 = sep(COV_ELLIPSOIDAL, 0.95)
    assert r50 < r90 < r95


def test_sep_scales_with_sigma() -> None:
    cov_scaled = 4.0 * COV_ELLIPSOIDAL
    assert sep_50(cov_scaled) == pytest.approx(2.0 * sep_50(COV_ELLIPSOIDAL), rel=1e-6)


# -- error handling --


def test_sep_invalid_conf() -> None:
    with pytest.raises(ValueError, match="conf"):
        sep(COV_SPHERICAL, 0.0)
    with pytest.raises(ValueError, match="conf"):
        sep(COV_SPHERICAL, 1.0)


def test_sep_invalid_shape() -> None:
    with pytest.raises(ValueError, match="shape"):
        sep(np.eye(2))


def test_sep_accepts_sequences() -> None:
    seq = ((25.0, 0.0, 0.0), (0.0, 25.0, 0.0), (0.0, 0.0, 25.0))
    assert sep_50(seq) == pytest.approx(sep_50(COV_SPHERICAL), rel=1e-10)


# -- extreme sigma ratios (i0e stability + axis choice) --


@pytest.mark.parametrize(
    "sigmas",
    [
        (1.0, 1.0, 20.0),
        (1.0, 1.0, 100.0),
        (20.0, 20.0, 1.0),
        (1.0, 5.0, 100.0),
    ],
    ids=["1:1:20", "1:1:100", "20:20:1", "1:5:100"],
)
def test_sep_extreme_ratio_monte_carlo(sigmas: tuple[float, float, float]) -> None:
    cov = np.diag([s**2 for s in sigmas])
    rng = np.random.default_rng(42)
    samples = rng.multivariate_normal([0, 0, 0], cov, size=2_000_000)
    radii = np.linalg.norm(samples, axis=1)
    mc_sep50 = float(np.median(radii))
    mc_sep95 = float(np.percentile(radii, 95))
    assert sep_50(cov) == pytest.approx(mc_sep50, rel=0.01)
    assert sep_95(cov) == pytest.approx(mc_sep95, rel=0.01)
