"""Average Scale Factor (ASF) for ellipse calibration analysis.

ASF and Median Scale Factor (MSF) measure how well reported uncertainty ellipses
match actual errors. An ASF/MSF of 1.0 indicates perfect calibration; values >1
indicate underestimated uncertainties (ellipses too small); values <1 indicate
overestimated uncertainties (ellipses too large).

The scale factor is computed by comparing empirical ellipse norms (Mahalanobis-like
distances normalized to 95% confidence) against the theoretical Rayleigh distribution
expected for properly calibrated 2D Gaussian uncertainties.
"""

from __future__ import annotations

import warnings
from typing import TYPE_CHECKING

import numpy as np
from scipy.stats import chi2
from scipy.stats import t as t_dist

from .quantile_ci import quantile_ci

if TYPE_CHECKING:
    from numpy.typing import NDArray

# Chi-squared value for 95% confidence with 2 DOF
# sqrt(chi2.ppf(0.95, 2)) ~= 2.4477
_CHI2_95_2DOF = float(np.sqrt(chi2.ppf(0.95, 2)))

# Minimum samples required for reliable ASF/MSF calculation
_MIN_SAMPLES = 25

# Floor value for ellipse norms to avoid division issues
_NORM_FLOOR = 0.005


def asf(
    ellipse_norms: NDArray[np.floating],
    confidence: float = 0.95,
    *,
    use_median: bool = False,
) -> tuple[float, float, float]:
    """Calculate Average Scale Factor (ASF) with confidence interval.

    Computes ASF or Median Scale Factor (MSF) that measures how well reported
    uncertainty ellipses match actual errors.

    The method:
        1. Compares empirical ellipse norms to theoretical Rayleigh distribution
        2. Scale Factor = ideal_norm / empirical_norm for each sample
        3. ASF = mean(SF), MSF = median(SF)
        4. CI computed using empirically-tuned formula (ASF) or quantile CI (MSF)

    Args:
        ellipse_norms: Array of ellipse norms (Mahalanobis-like distances normalized
            to 95% confidence ellipse). NaN values are removed with a warning.
        confidence: Confidence level for the interval, in (0, 1). Default 0.95.
        use_median: If True, compute Median Scale Factor instead of Average.
            Default False (compute ASF).

    Returns:
        Tuple of (scale_factor, ci_lower, ci_upper). Returns (nan, nan, nan) if
        fewer than 25 valid samples are provided.

    Example:
        >>> norms = np.random.rayleigh(scale=1.0, size=100) / _CHI2_95_2DOF
        >>> asf_val, lower, upper = asf(norms)
        >>> msf_val, lower, upper = asf(norms, use_median=True)
    """
    ellipse_norms = np.asarray(ellipse_norms).ravel()

    # Remove NaN values with warning
    nan_mask = np.isnan(ellipse_norms)
    if np.any(nan_mask):
        warnings.warn(
            f"Removing {np.sum(nan_mask)} NaN ellipse norms",
            UserWarning,
            stacklevel=2,
        )
        ellipse_norms = ellipse_norms[~nan_mask]

    n = len(ellipse_norms)
    if n < _MIN_SAMPLES:
        return np.nan, np.nan, np.nan

    # Sort and apply floor to avoid division by very small values
    empirical = np.sort(ellipse_norms)
    empirical = np.maximum(empirical, _NORM_FLOOR)

    # Percentiles: (1/n, 2/n, ..., n/n) offset by 1/(2n) for better estimation
    percentiles = (np.arange(1, n + 1) / n) - (1 / (2 * n))

    # Ideal values from chi-squared distribution (Rayleigh for 2D Gaussian)
    # x_ideal = sqrt(chi2.ppf(percentile, 2)) / sqrt(chi2.ppf(0.95, 2))
    # chi2.ppf(p, 2) == -2 * ln(1 - p)
    ideal = np.sqrt(-2 * np.log(1 - percentiles)) / _CHI2_95_2DOF

    # Scale factors: how much to scale ellipse to match ideal
    sf = ideal / empirical

    if not use_median:
        # Average Scale Factor
        asf_val = float(np.mean(sf))

        # CI using empirically-tuned formula
        # Combines standard error of mean with additional variance term
        t_crit = float(t_dist.ppf(1 - (1 - confidence) / 2, n))
        sem_term = (np.std(sf, ddof=1) / np.sqrt(n)) ** 2
        asf_term = (asf_val**2 / np.sqrt(2 * n)) ** 2
        std_est = np.sqrt(sem_term + asf_term)
        ci_lower = asf_val - std_est * t_crit
        ci_upper = asf_val + std_est * t_crit
    else:
        # Median Scale Factor
        asf_val = float(np.median(sf))
        ci_lower, ci_upper = quantile_ci(sf, quantile=0.5, confidence=confidence)
        ci_lower = float(ci_lower)
        ci_upper = float(ci_upper)

    return asf_val, ci_lower, ci_upper
