"""
Verification: does the L0 closure posterior recover the truth PDF (LH_PARAM_20250519)?

Compares posterior PDF curves (from posterior samples + pdf_model.pkl) against the
truth PDF evaluated from the LHAPDF grid via validphys, in the evolution basis,
for the independently fitted flavours: Sigma, g, V, V3.
"""
import dill
import numpy as np
import pandas as pd
import jax
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

from validphys.loader import Loader
from validphys import convolution

RUN = "/Users/mukesh/colibri-runs/lh_fit_closure_test"
OUT = ("/Users/mukesh/Library/CloudStorage/GoogleDrive-mukesh@merakilabs.com"
       "/My Drive/MB-Git-Repo/AI Test/53-Proton-PDF/output/lh_fit_closure_test/verification")
import os

os.makedirs(OUT, exist_ok=True)

Q0 = 1.65
# x-grid for comparison; data region for SLAC+BCDMS DIS is roughly x ~ 0.01-0.75
xgrid = np.logspace(-3, np.log10(0.95), 120)

# --- posterior curves -------------------------------------------------------
with open(f"{RUN}/pdf_model.pkl", "rb") as f:
    model = dill.load(f)

samples = pd.read_csv(f"{RUN}/full_posterior_sample.csv", index_col=0)
assert list(samples.columns) == list(model.param_names), (
    samples.columns, model.param_names)

pdf_func = model.grid_values_func(xgrid)
curves = np.array(jax.vmap(lambda p: pdf_func(p))(samples.values))  # (Ns, 14, Nx)
print(f"posterior samples: {curves.shape[0]}, flavours: {curves.shape[1]}, x-points: {curves.shape[2]}")

post_mean = curves.mean(axis=0)
post_std = curves.std(axis=0)

# --- truth curves -----------------------------------------------------------
pdf = Loader().check_pdf("LH_PARAM_20250519")
truth_grid = convolution.evolution.grid_values(
    pdf, convolution.FK_FLAVOURS, xgrid, [Q0]
)  # (Nmem, 14, Nx, 1)
truth = np.asarray(truth_grid)[0, :, :, 0]  # central member

# --- compare fitted flavours ------------------------------------------------
FLAVOURS = {1: r"$\Sigma$", 2: "g", 3: "V", 4: "V3"}
rows = []
fig, axes = plt.subplots(2, 2, figsize=(11, 8))
for ax, (idx, name) in zip(axes.flat, FLAVOURS.items()):
    m, s, t = post_mean[idx], post_std[idx], truth[idx]
    pull = (m - t) / np.where(s > 0, s, np.nan)
    cover68 = np.mean(np.abs(pull) <= 1.0)
    reldev = np.abs(m - t) / np.maximum(np.abs(t), 1e-12)
    rows.append({
        "flavour": name, "max_abs_pull": np.nanmax(np.abs(pull)),
        "mean_abs_pull": np.nanmean(np.abs(pull)),
        "coverage_within_1sigma": cover68,
        "median_rel_dev_of_mean": np.median(reldev),
    })
    ax.fill_between(xgrid, m - s, m + s, alpha=0.35, color="#0e7c7b",
                    label="posterior 68%")
    ax.plot(xgrid, m, color="#0e7c7b", lw=1.2, label="posterior mean")
    ax.plot(xgrid, t, "k--", lw=1.4, label="truth (LH_PARAM_20250519)")
    ax.set_xscale("log")
    ax.set_title(f"{name}  (evolution basis, Q0 = {Q0} GeV)")
    ax.set_xlabel("x")
    ax.axvspan(1e-3, 0.01, color="gray", alpha=0.12)
    if idx == 1:
        ax.legend(fontsize=8)
fig.suptitle("L0 closure: posterior vs truth — shaded x<0.01 is outside DIS data coverage")
fig.tight_layout()
fig.savefig(f"{OUT}/posterior_vs_truth.png", dpi=160)

report = pd.DataFrame(rows)
report.to_csv(f"{OUT}/pull_summary.csv", index=False)
print(report.to_string(index=False))

# restrict to data-covered region x > 0.01
mask = xgrid > 0.01
print("\n--- restricted to x > 0.01 (data region) ---")
for idx, name in FLAVOURS.items():
    m, s, t = post_mean[idx][mask], post_std[idx][mask], truth[idx][mask]
    pull = np.abs((m - t) / np.where(s > 0, s, np.nan))
    print(f"{name:8s} max|pull|={np.nanmax(pull):7.3f}  mean|pull|={np.nanmean(pull):6.3f}  "
          f"1sigma-coverage={np.mean(pull<=1):5.1%}")