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

Shows the distribution of ellipse norms compared to the expected Rayleigh
distribution for properly calibrated uncertainties.

- Distribution shifted right: Uncertainties underestimated (ASF > 1)
- Distribution matches theoretical: Well-calibrated (ASF ≈ 1)
- Distribution shifted left: Uncertainties overestimated (ASF < 1)
"""

from __future__ import annotations

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

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.

    Returns:
        Array of ellipse norms normalized to 95% confidence.
    """
    errors = RNG.standard_normal((n, 2)) * scale
    mahal = np.sqrt(np.sum(errors**2, axis=1))
    return mahal / CHI2_95_2DOF


def histogram_trace(
    data: np.ndarray,
    bins: np.ndarray,
    name: str,
) -> tuple[np.ndarray, np.ndarray, dict]:
    """Create histogram trace data for scatter plot.

    Args:
        data: Array of values to histogram.
        bins: Bin edges.
        name: Trace name for legend.

    Returns:
        Tuple of (bin_centers, counts, trace_args).
    """
    counts, _ = np.histogram(data, bins=bins, density=True)
    centers = (bins[:-1] + bins[1:]) / 2
    style = {"mode": "lines+markers", "marker": {"size": 5}}
    return centers, counts, {"name": name, **style}


def main() -> None:
    """Generate histogram comparing empirical norms to theoretical Rayleigh."""
    n = 2000
    bins = np.linspace(0, 2.5, 100)

    # Generate samples with different calibrations
    norms_under = generate_ellipse_norms(n, scale=1.4)  # ASF > 1
    norms_good = generate_ellipse_norms(n, scale=1.0)  # ASF ≈ 1

    # Theoretical Rayleigh PDF (normalized to 95% confidence)
    x_theory = np.linspace(0.01, 2.5, 200)
    # Rayleigh with scale=1, then normalize by CHI2_95_2DOF
    pdf_theory = rayleigh.pdf(x_theory * CHI2_95_2DOF) * CHI2_95_2DOF

    # Build histogram traces
    h_under = histogram_trace(norms_under, bins, "Underestimated (ASF>1)")
    h_good = histogram_trace(norms_good, bins, "Well-calibrated (ASF≈1)")

    scatter(
        (x_theory, pdf_theory, {"name": "Theoretical Rayleigh", "line": {"width": 3}}),
        h_under,
        h_good,
        title="Scale Factor Histogram",
        xaxis_title="Ellipse Norm (normalized to 95%)",
        yaxis_title="Probability Density",
    )


if __name__ == "__main__":
    main()
