"""Spherical Error Probable (SEP) from a 3D covariance matrix.

SEP is the radius of a sphere centered at the mean that contains a given
fraction of the probability mass of a 3D Gaussian distribution. For a
spherical distribution (equal eigenvalues) the radial distance follows a
Maxwell-Boltzmann distribution and the solution is closed-form via the
chi distribution with 3 degrees of freedom. For a general ellipsoidal
distribution (unequal eigenvalues) the CDF is computed via numerical
integration and inverted with root-finding.

Numerical approach:
    The 3D Gaussian is integrated in spherical coordinates. The azimuthal
    (theta) integral evaluates analytically to a modified Bessel function
    I0, leaving a 1D integral over the polar angle phi (evaluated with
    fixed-order Gauss-Legendre quadrature) nested inside the radial
    integral (evaluated with adaptive quadrature). Two techniques ensure
    correctness for arbitrary sigma ratios, including extreme cases like
    1:1:100 or 20:20:1:

    1. **Exponentially-scaled Bessel function**: The identity
       ``exp(-a) * I0(b) = exp(-(a-b)) * i0e(b)`` avoids I0 overflow
       when the Bessel argument is large.

    2. **Optimal axis assignment**: The three principal axes are assigned
       to the spherical coordinate system so that the two most similar
       sigmas (in inverse-square sense) form the theta-plane pair. This
       minimizes the I0 argument, keeping the phi integrand smooth enough
       for Gauss-Legendre quadrature. Without this, pairing two very
       different sigmas (e.g. 100 and 1) causes the I0 argument to
       nearly cancel the exponential decay, producing a sharply peaked
       integrand that fixed-order quadrature cannot resolve.

All functions accept a 3x3 covariance matrix. Units of the result match
the units of the covariance input (e.g. meters if covariance is in m^2).
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
from scipy.integrate import quad
from scipy.optimize import brentq
from scipy.special import i0e
from scipy.stats import chi, chi2

from gri_utils.matrices import decompose

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray

_REL_TOL = 1e-12

# Gauss-Legendre nodes and weights for the phi integral over [0, pi].
# 192 points provides sub-0.001% accuracy even for extreme sigma ratios
# (e.g. 20:20:1) where the phi integrand becomes sharply peaked.
_GL_ORDER = 192
_gl_nodes_01, _gl_weights_01 = np.polynomial.legendre.leggauss(_GL_ORDER)
# Map from [-1, 1] to [0, pi]
_GL_PHI = np.pi / 2.0 * (_gl_nodes_01 + 1.0)
_GL_W = np.pi / 2.0 * _gl_weights_01
# Pre-compute trig values at the quadrature nodes
_GL_SIN = np.sin(_GL_PHI)
_GL_SIN2 = _GL_SIN**2
_GL_COS2 = np.cos(_GL_PHI) ** 2
_GL_SIN_W = _GL_SIN * _GL_W


def sep(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
    conf: float = 0.50,
) -> float:
    """Spherical Error Probable at a given confidence level.

    Computes the radius of a sphere centered at the mean that contains
    ``conf`` fraction of the probability mass of the 3D Gaussian
    described by ``cov``.

    Args:
        cov: 3x3 covariance matrix (1-sigma, units^2).
        conf: Confidence level in (0, 1). Default 0.50 (50th percentile).

    Returns:
        SEP radius in the same length units as the covariance.

    Raises:
        ValueError: If conf is not in (0, 1) or covariance shape is invalid.
    """
    if not 0 < conf < 1:
        msg = f"conf must be in (0, 1), got {conf}"
        raise ValueError(msg)
    sig1, sig2, sig3 = _extract_sigmas_3d(cov)

    # Upper bound for root-finding: use 1-1e-8 quantile for extreme ratios
    r_upper = sig1 * np.sqrt(chi2.ppf(1 - 1e-8, 3))

    return float(
        brentq(
            lambda r: _cdf(r, sig1, sig2, sig3) - conf,
            0,
            r_upper,
        ),
    )


def sep_sigma(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
    sigma: float = 1.0,
) -> float:
    """SEP at a given number of standard deviations.

    Converts sigma to a confidence level via the 3D chi-squared distribution,
    then computes SEP at that confidence.

    Args:
        cov: 3x3 covariance matrix (1-sigma, units^2).
        sigma: Number of standard deviations. Default 1.0.

    Returns:
        SEP radius in the same length units as the covariance.
    """
    conf = float(chi2.cdf(sigma**2, df=3))
    return sep(cov, conf)


def sep_50(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> float:
    """SEP at 50% confidence.

    Args:
        cov: 3x3 covariance matrix (1-sigma, units^2).

    Returns:
        SEP radius in the same length units as the covariance.
    """
    return sep(cov, 0.50)


def sep_90(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> float:
    """SEP at 90% confidence.

    Args:
        cov: 3x3 covariance matrix (1-sigma, units^2).

    Returns:
        SEP radius in the same length units as the covariance.
    """
    return sep(cov, 0.90)


def sep_95(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> float:
    """SEP at 95% confidence.

    Args:
        cov: 3x3 covariance matrix (1-sigma, units^2).

    Returns:
        SEP radius in the same length units as the covariance.
    """
    return sep(cov, 0.95)


def sep_2sigma(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> float:
    """SEP at 2-sigma confidence.

    Args:
        cov: 3x3 covariance matrix (1-sigma, units^2).

    Returns:
        SEP radius in the same length units as the covariance.
    """
    return sep_sigma(cov, 2.0)


# ---------------------------------------------------------------------------
# Private helpers
# ---------------------------------------------------------------------------


def _extract_sigmas_3d(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> tuple[float, float, float]:
    """Extract three standard deviations (sqrt of eigenvalues) from a 3x3 cov.

    Args:
        cov: 3x3 covariance matrix.

    Returns:
        Tuple of (sig1, sig2, sig3) sorted largest to smallest.

    Raises:
        ValueError: If shape is not (3, 3).
    """
    cov = np.asarray(cov, dtype=np.float64)
    if cov.shape != (3, 3):
        msg = f"Covariance shape must be (3, 3), got {cov.shape}"
        raise ValueError(msg)
    vals, _ = decompose(cov)
    return float(vals[0] ** 0.5), float(vals[1] ** 0.5), float(vals[2] ** 0.5)


def _choose_axes(
    sig1: float,
    sig2: float,
    sig3: float,
) -> tuple[float, float, float]:
    """Choose optimal axis assignment for the spherical coordinate integral.

    The 3D radial PDF is derived by integrating the Gaussian in spherical
    coordinates. The azimuthal (theta) integral yields a modified Bessel
    function I0 whose argument depends on how different the two theta-plane
    sigmas are. We assign the two most similar sigmas to the theta plane
    to minimize the I0 argument, ensuring both numerical stability and
    fast Gauss-Legendre convergence for the polar (phi) integral.

    Args:
        sig1: First principal standard deviation.
        sig2: Second principal standard deviation.
        sig3: Third principal standard deviation.

    Returns:
        Tuple (sa, sb, sc) where sa, sb are the theta-plane pair
        (sa >= sb) and sc is the polar axis.
    """
    sigs = sorted([sig1, sig2, sig3])
    s1, s2, s3 = sigs  # ascending

    def diff_ratio(a: float, b: float) -> float:
        ia, ib = 1.0 / a**2, 1.0 / b**2
        return abs(ia - ib) / (ia + ib)

    candidates = [
        (max(s1, s2), min(s1, s2), s3, diff_ratio(s1, s2)),
        (max(s1, s3), min(s1, s3), s2, diff_ratio(s1, s3)),
        (max(s2, s3), min(s2, s3), s1, diff_ratio(s2, s3)),
    ]
    best = min(candidates, key=lambda x: x[3])
    return best[0], best[1], best[2]


def _cdf(
    r: float,
    sig1: float,
    sig2: float,
    sig3: float,
) -> float:
    """CDF of the radial distance for a zero-mean trivariate normal.

    For a 3D Gaussian with principal standard deviations sig1, sig2, sig3,
    compute P(R <= r) where R = sqrt(X^2 + Y^2 + Z^2).

    When all sigmas are equal this is the chi distribution CDF with k=3.
    Otherwise integrates the generalized radial density numerically.

    The PDF of R is derived by integrating the 3D Gaussian in spherical
    coordinates. The theta integral evaluates analytically to a modified
    Bessel function I0, leaving a single integral over the polar angle phi:

        f(r) = (r^2 / (sqrt(2*pi) * sa * sb * sc))
               * integral_{0}^{pi}
                 sin(phi) * exp(-r^2 * A(phi))
                 * I0(r^2 * B(phi))
                 d_phi

    where A(phi) and B(phi) depend on the three sigmas and the polar
    angle. The CDF is then obtained by integrating f(r) from 0 to the
    target radius.

    To avoid I0 overflow for extreme sigma ratios, the identity
    exp(-a)*I0(b) = exp(-(a-b))*i0e(b) is used, where i0e is the
    exponentially scaled Bessel function.

    The axis assignment (sa, sb for theta, sc for polar) is chosen to
    minimize the I0 argument, ensuring numerical stability and accurate
    Gauss-Legendre quadrature for all sigma ratios.

    The phi integral is evaluated using fixed-order Gauss-Legendre
    quadrature with pre-computed nodes and weights for speed.

    Args:
        r: Radius at which to evaluate the CDF.
        sig1: Standard deviation along the first principal axis.
        sig2: Standard deviation along the second principal axis.
        sig3: Standard deviation along the third principal axis.

    Returns:
        Probability that the radial distance is <= r.
    """
    if r <= 0:
        return 0.0

    # Spherical case: chi distribution with k=3
    max_sig = max(sig1, sig2, sig3)
    if abs(sig1 - sig2) / max_sig < _REL_TOL and abs(sig2 - sig3) / max_sig < _REL_TOL:
        return float(chi.cdf(r / sig1, df=3))

    # Choose optimal axis assignment
    sa, sb, sc = _choose_axes(sig1, sig2, sig3)

    inv_sa_sq = 1.0 / sa**2
    inv_sb_sq = 1.0 / sb**2
    inv_sc_sq = 1.0 / sc**2
    sum_ab = (inv_sa_sq + inv_sb_sq) / 4.0
    diff_ab = abs(inv_sa_sq - inv_sb_sq) / 4.0
    norm = 1.0 / (np.sqrt(2.0 * np.pi) * sa * sb * sc)

    def radial_pdf(rr: float) -> float:
        """PDF of radial distance at rr, with phi via Gauss-Legendre."""
        rr_sq = rr**2
        a_vals = rr_sq * (_GL_SIN2 * sum_ab + _GL_COS2 * inv_sc_sq / 2.0)
        b_vals = rr_sq * _GL_SIN2 * diff_ab
        # exp(-a) * I0(b) = exp(-(a - b)) * i0e(b)
        angular = float(
            np.sum(_GL_SIN_W * np.exp(-(a_vals - b_vals)) * i0e(b_vals)),
        )
        return rr_sq * norm * angular

    result, _ = quad(radial_pdf, 0, r, limit=100)
    return float(result)
