"""
Comparison of density estimators on real model output scores.

For each class (negative / positive) plots four curves against the empirical CDF:
  - ECDF              — ground truth step function
  - KDE (ISJ)         — Gaussian-kernel KDE with Improved Sheather-Jones bandwidth
  - Gaussian fit      — MLE normal distribution
  - Beta fit          — sigmoid-transformed scores fitted with Beta MLE, CDF
                        mapped back to the original logit axis

Data: test split of a ROLL model trained on the glass0 dataset.

Run from the impl directory (KDEpy lives in that flake):
    cd impl && nix develop --command python3 ../thesis/content/method/figures/density_fit_comparison.py
"""

import sys, pickle
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.stats import norm, beta as sp_beta
from scipy.special import expit
from KDEpy import FFTKDE
from KDEpy.bw_selection import silvermans_rule, improved_sheather_jones

# Resolve paths relative to this script's location so the script runs from anywhere.
_HERE    = Path(__file__).resolve().parent          # .../thesis/content/method/figures
_THESIS  = _HERE.parents[2]                         # .../thesis
_IMPL    = _THESIS.parent / "impl"                  # .../impl

sys.path.insert(0, str(_IMPL))  # pickle needs 'import src.experiment' → impl/src/experiment.py

# ── Config ──────────────────────────────────────────────────────────────────
RESULT_PKL = _IMPL / "results/glass0/2026-07-08-22-25/roll/0/test-res.pkl"
OUT_DIR    = _HERE
SPLIT      = "test"

C_NEG = "#4477AA"
C_POS = "#CC6633"

plt.rcParams.update({
    "text.usetex":      False,
    "mathtext.fontset": "cm",
    "pdf.fonttype":     42,
})

# ── Load data ────────────────────────────────────────────────────────────────
with open(RESULT_PKL, "rb") as f:
    ep = pickle.load(f)

res = ep.split_results[SPLIT]
yh  = res.yh.astype(np.float64)   # raw logit scores
y   = res.y

scores_neg = yh[y == 0]
scores_pos = yh[y == 1]

print(f"Loaded {SPLIT} split: {len(scores_neg)} negatives, {len(scores_pos)} positives")
print(f"Score range: [{yh.min():.3f}, {yh.max():.3f}]")


# ── Helpers ──────────────────────────────────────────────────────────────────
def ecdf(scores):
    """Return (x, p) for a step-function ECDF."""
    s = np.sort(scores)
    p = np.arange(1, len(s) + 1) / len(s)
    return s, p


def _kde_pdf(scores, x_grid, h):
    pdf = FFTKDE(kernel="gaussian", bw=h).fit(scores).evaluate(x_grid)
    dx  = x_grid[1] - x_grid[0]
    cdf = np.cumsum(pdf) * dx
    cdf /= cdf[-1]
    return cdf

def kde_isj_cdf(scores, x_grid):
    """ISJ bandwidth; Silverman fallback on failure."""
    data = scores[:, np.newaxis].astype(np.float64)
    try:
        h = float(improved_sheather_jones(data))
    except Exception:
        h = float(silvermans_rule(data))
    return _kde_pdf(scores, x_grid, h)

def kde_silverman_cdf(scores, x_grid):
    h = float(silvermans_rule(scores[:, np.newaxis].astype(np.float64)))
    return _kde_pdf(scores, x_grid, h)


def gauss_cdf(scores, x_grid):
    loc, scale = norm.fit(scores)
    return norm.cdf(x_grid, loc, scale)


def beta_cdf(scores, x_grid):
    """Fit Beta in sigmoid space; evaluate CDF back on the logit x_grid."""
    s01 = expit(scores).clip(1e-6, 1 - 1e-6)
    a, b, _, _ = sp_beta.fit(s01, floc=0, fscale=1)
    return sp_beta.cdf(expit(x_grid), a, b)


# ── Shared x-grid (covers both classes with a small margin) ──────────────────
margin = 0.5
x_lo = yh.min() - margin
x_hi = yh.max() + margin
x_grid = np.linspace(x_lo, x_hi, 2048)


# ── Compute all curves ───────────────────────────────────────────────────────
neg_ecdf_x, neg_ecdf_p   = ecdf(scores_neg)
neg_cdf_kde              = kde_isj_cdf(scores_neg, x_grid)
neg_cdf_silverman        = kde_silverman_cdf(scores_neg, x_grid)
neg_cdf_gauss            = gauss_cdf(scores_neg, x_grid)
neg_cdf_beta             = beta_cdf(scores_neg, x_grid)

pos_ecdf_x, pos_ecdf_p   = ecdf(scores_pos)
pos_cdf_kde              = kde_isj_cdf(scores_pos, x_grid)
pos_cdf_silverman        = kde_silverman_cdf(scores_pos, x_grid)
pos_cdf_gauss            = gauss_cdf(scores_pos, x_grid)
pos_cdf_beta             = beta_cdf(scores_pos, x_grid)


# ── Plot ─────────────────────────────────────────────────────────────────────
fig, (ax0, ax1) = plt.subplots(
    2, 1, figsize=(5.5, 4.8), sharex=True,
    gridspec_kw={"height_ratios": [1, 1]},
)
fig.subplots_adjust(hspace=0.08)


def _plot_class(ax, color, ecdf_x, ecdf_p,
                cdf_isj, cdf_silverman, cdf_gauss, cdf_beta, label_prefix):
    ax.step(ecdf_x, ecdf_p,
            color=color, lw=1.0, alpha=0.55, where="post",
            label="ECDF")
    ax.plot(x_grid, cdf_isj,
            color=color, lw=1.8, linestyle="-",
            label="KDE (ISJ)")
    ax.plot(x_grid, cdf_silverman,
            color=color, lw=1.4, linestyle="-.",
            label="KDE (Silverman)")
    ax.plot(x_grid, cdf_gauss,
            color=color, lw=1.4, linestyle="--",
            label="Gaussian")
    ax.plot(x_grid, cdf_beta,
            color=color, lw=1.6, linestyle=":",
            label="Beta")
    ax.set_ylabel("CDF", fontsize=9)
    ax.set_yticks([0.0, 0.25, 0.50, 0.75, 1.0])
    ax.set_yticklabels(["0", ".25", ".5", ".75", "1"], fontsize=7.5)
    ax.tick_params(labelsize=8)
    ax.spines[["top", "right"]].set_visible(False)
    ax.text(0.97, 0.05, label_prefix,
            transform=ax.transAxes, fontsize=8,
            ha="right", va="bottom", color=color)


_plot_class(ax0, C_NEG,
            neg_ecdf_x, neg_ecdf_p,
            neg_cdf_kde, neg_cdf_silverman, neg_cdf_gauss, neg_cdf_beta,
            r"Negative class ($y=0$)")

_plot_class(ax1, C_POS,
            pos_ecdf_x, pos_ecdf_p,
            pos_cdf_kde, pos_cdf_silverman, pos_cdf_gauss, pos_cdf_beta,
            r"Positive class ($y=1$)")

ax0.legend(fontsize=7.5, loc="upper left", framealpha=0.88)
ax1.set_xlabel(r"Score $f_\theta(\mathbf{x})$", fontsize=9)
ax0.set_xlim(x_lo, x_hi)

for path_suffix in ("pdf", "png"):
    kw = {"bbox_inches": "tight"}
    if path_suffix == "png":
        kw["dpi"] = 180
    fig.savefig(OUT_DIR / f"density_fit_comparison.{path_suffix}", **kw)

print(f"Saved to {OUT_DIR}/density_fit_comparison.{{pdf,png}}")
