"""Demo of scale factor Q-Q plot - empirical vs theoretical Rayleigh distribution.

A Q-Q plot visualizes how well uncertainty ellipses are calibrated by comparing
the sorted empirical ellipse norms against the theoretical Rayleigh distribution.

- Points on the diagonal (y=x): Perfect calibration (ASF = 1.0)
- Points above diagonal: Uncertainties underestimated (ellipses too small)
- Points below diagonal: Uncertainties overestimated (ellipses too large)
"""

from __future__ import annotations

import numpy as np
from gri_plot import scatter
from scipy.stats import chi2

RNG = np.random.default_rng(42)

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


def generate_ellipse_norms(n: int, scale: float = 1.0) -> np.ndarray:
    """Generate ellipse norms from 2D Gaussian errors.

    Args:
        n: Number of samples.
        scale: Scale factor applied to errors. 1.0 = perfect calibration,
               >1 = underestimated uncertainty, <1 = overestimated.

    Returns:
        Array of ellipse norms (Mahalanobis-like distances normalized to 95%).
    """
    # 2D Gaussian errors scaled by the scale factor
    errors = RNG.standard_normal((n, 2)) * scale
    # Mahalanobis distance for unit covariance is just the Euclidean norm
    mahal = np.sqrt(np.sum(errors**2, axis=1))
    # Normalize to 95% confidence ellipse
    return mahal / CHI2_95_2DOF


def theoretical_quantiles(n: int) -> np.ndarray:
    """Compute theoretical Rayleigh quantiles for Q-Q plot.

    Args:
        n: Number of samples.

    Returns:
        Theoretical quantile values from Rayleigh distribution.
    """
    # Percentiles offset by 1/(2n) for better estimation
    percentiles = (np.arange(1, n + 1) / n) - (1 / (2 * n))
    # Rayleigh quantiles: sqrt(-2*ln(1-p)) normalized to 95% confidence
    return np.sqrt(-2 * np.log(1 - percentiles)) / CHI2_95_2DOF


def main() -> None:
    """Generate Q-Q plots for different calibration scenarios."""
    n = 200

    # Three scenarios: underestimated, calibrated, overestimated
    norms_under = np.sort(generate_ellipse_norms(n, scale=1.5))  # ASF > 1
    norms_good = np.sort(generate_ellipse_norms(n, scale=1.0))  # ASF ≈ 1
    norms_over = np.sort(generate_ellipse_norms(n, scale=0.7))  # ASF < 1

    theoretical = theoretical_quantiles(n)

    # Reference line (perfect calibration)
    ref_line = np.array([0, theoretical[-1]])

    marker = {"mode": "markers", "marker": {"size": 4}}
    scatter(
        (theoretical, norms_under, {"name": "Underestimated (ASF>1)", **marker}),
        (theoretical, norms_good, {"name": "Well-calibrated (ASF≈1)", **marker}),
        (theoretical, norms_over, {"name": "Overestimated (ASF<1)", **marker}),
        (ref_line, ref_line, {"name": "Perfect (y=x)", "line": {"dash": "dash"}}),
        title="Scale Factor Q-Q Plot",
        xaxis_title="Theoretical Quantile (Rayleigh)",
        yaxis_title="Empirical Ellipse Norm",
    )


if __name__ == "__main__":
    main()
