"""Circular Error Probable (CEP) from a 2D covariance matrix.

CEP is the radius of a circle centered at the mean that contains a given
fraction of the probability mass of a 2D Gaussian distribution. For a
circular distribution (equal eigenvalues) the radial distance follows a
Rayleigh distribution and the solution is closed-form. For an elliptical
distribution (unequal eigenvalues) the CDF has no closed form and is
computed via numerical integration (Beckmann distribution), then inverted
with root-finding.

Numerical approach:
    The Beckmann PDF involves the modified Bessel function I0, which
    overflows for extreme sigma ratios (e.g. 1:100 or larger). This is
    avoided by using the exponentially-scaled Bessel function i0e via
    the identity ``exp(-a) * I0(b) = exp(-(a-b)) * i0e(b)``, which is
    numerically stable for all sigma ratios. The CDF integral is
    evaluated with adaptive quadrature (scipy.integrate.quad) and the
    quantile is found via Brent root-finding.

All functions accept a 2x2 covariance matrix (or the 2x2 projection of a
3x3). 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 chi2

from gri_utils.matrices import decompose

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray

_REL_TOL = 1e-12


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

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

    Args:
        cov: 2x2 or 3x3 covariance matrix (1-sigma, units^2). If 3x3,
            the upper-left 2x2 block is used.
        conf: Confidence level in (0, 1). Default 0.50 (50th percentile).

    Returns:
        CEP 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 = _extract_sigmas(cov)

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

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


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

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

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

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


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

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

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


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

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

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


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

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

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


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

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

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


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


def _extract_sigmas(
    cov: NDArray[np.floating] | Sequence[Sequence[float | int]],
) -> tuple[float, float]:
    """Extract the two standard deviations (sqrt of eigenvalues) from a 2x2 cov.

    Accepts 2x2 or 3x3 (projects to 2x2 via upper-left block).

    Args:
        cov: 2x2 or 3x3 covariance matrix.

    Returns:
        Tuple of (sigma_major, sigma_minor) with sigma_major >= sigma_minor.

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


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

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

    When sig1 == sig2 this is the Rayleigh CDF. Otherwise it integrates
    the Beckmann (Hoyt) density numerically.

    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.

    Returns:
        Probability that the radial distance is <= r.
    """
    if r <= 0:
        return 0.0
    # Circular case: Rayleigh CDF = 1 - exp(-r^2 / (2 * sigma^2))
    if abs(sig1 - sig2) / max(sig1, sig2) < _REL_TOL:
        return float(1.0 - np.exp(-(r**2) / (2.0 * sig1**2)))

    # General case: integrate the Beckmann/Hoyt PDF
    # PDF(r) = (r / (sig1 * sig2))
    #          * exp(-r^2 * (1/sig1^2 + 1/sig2^2) / 4)
    #          * I0(r^2 * (1/sig1^2 - 1/sig2^2) / 4)
    # where I0 is the modified Bessel function of the first kind, order 0.
    #
    # To avoid overflow for extreme sigma ratios, use the identity:
    #   exp(-a) * I0(b) = exp(-(a - b)) * i0e(b)
    # where i0e(x) = exp(-|x|) * I0(x) is the exponentially scaled Bessel.
    # Since a = sum_inv >= diff_inv = b always, (a - b) >= 0 and this is stable.
    inv_s1_sq = 1.0 / sig1**2
    inv_s2_sq = 1.0 / sig2**2
    sum_inv = (inv_s1_sq + inv_s2_sq) / 4.0
    diff_inv = abs(inv_s1_sq - inv_s2_sq) / 4.0
    net_decay = sum_inv - diff_inv
    norm = 1.0 / (sig1 * sig2)

    def pdf(rr: float) -> float:
        rr_sq = rr**2
        return float(
            rr * norm * np.exp(-rr_sq * net_decay) * i0e(rr_sq * diff_inv),
        )

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