"""
Cross-method comparison: Bayesian vs MC-replica vs Hessian closure fits (L0 & L1)
at the PDF-curve level, against the truth PDF LH_PARAM_20250519.
"""
import json
import os
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

RUNS = "/Users/mukesh/colibri-runs"
OUT = ("/Users/mukesh/Library/CloudStorage/GoogleDrive-mukesh@merakilabs.com"
       "/My Drive/MB-Git-Repo/AI Test/53-Proton-PDF/output/crosscheck_methods")
os.makedirs(OUT, exist_ok=True)

Q0 = 1.65
xgrid = np.logspace(-3, np.log10(0.95), 120)

with open(f"{RUNS}/lh_fit_closure_test/pdf_model.pkl", "rb") as f:
    model = dill.load(f)
pdf_func = model.grid_values_func(xgrid)
curves_of = lambda P: np.array(jax.vmap(lambda p: pdf_func(p))(np.atleast_2d(P)))

pdf = Loader().check_pdf("LH_PARAM_20250519")
truth = np.asarray(convolution.evolution.grid_values(
    pdf, convolution.FK_FLAVOURS, xgrid, [Q0]))[0, :, :, 0]

FLAV = {1: r"$\Sigma$", 2: "g", 3: "V"}
COLORS = {"Bayesian": "#0e7c7b", "MC replicas": "#d96c2c", "Hessian": "#5b4a9f"}

def band(curves):
    return curves.mean(0), curves.std(0)

def hessian_band(run):
    d = json.load(open(f"{run}/hessian_fit_summary.json"))
    central = curves_of(np.array(d["optimized_parameters"]))[0]
    mem = pd.read_csv(f"{run}/hessian_result.csv", index_col=0).values
    mc = curves_of(mem)  # (26, 14, Nx): ± pairs per eigenvector
    plus, minus = mc[0::2], mc[1::2]
    sig = np.sqrt(np.nansum(((plus - minus) / 2.0) ** 2, axis=0))
    return central, sig

summary_rows = []
for lvl, bay_dir, mc_dir, hess_dir in [
    ("L0", "lh_fit_closure_test", "lh_mc_L0", "lh_hessian_L0"),
    ("L1", "lh_fit_closure_test_L1", "lh_mc_L1", "lh_hessian_L1"),
]:
    bay = curves_of(pd.read_csv(f"{RUNS}/{bay_dir}/full_posterior_sample.csv", index_col=0).values)
    mc = curves_of(pd.read_csv(f"{RUNS}/{mc_dir}/mc_result.csv", index_col=0).values)
    hc, hs = hessian_band(f"{RUNS}/{hess_dir}")
    bands = {"Bayesian": band(bay), "MC replicas": band(mc), "Hessian": (hc, hs)}

    fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.2))
    for ax, (idx, name) in zip(axes, FLAV.items()):
        for meth, (m, s) in bands.items():
            ax.fill_between(xgrid, m[idx] - s[idx], m[idx] + s[idx],
                            alpha=0.25, color=COLORS[meth])
            ax.plot(xgrid, m[idx], color=COLORS[meth], lw=1.3, label=meth)
        ax.plot(xgrid, truth[idx], "k--", lw=1.5, label="truth")
        ax.set_xscale("log"); ax.set_xlabel("x")
        ax.set_title(f"{name} — {lvl}")
        if idx == 3:
            lo = min(truth[idx].min(), 0) - 3
            ax.set_ylim(lo, truth[idx].max() + 3)  # cap runaway Hessian band
        if idx == 1:
            ax.legend(fontsize=8)
    fig.suptitle(f"{lvl} closure: three inference methods vs truth (bands = 1σ; V panel y-clipped)")
    fig.tight_layout()
    fig.savefig(f"{OUT}/methods_{lvl}.png", dpi=160)

    # numeric summary on well-constrained flavours, data region x>0.01
    mask = xgrid > 0.01
    for idx in (1, 2):
        t = truth[idx][mask]
        for meth, (m, s) in bands.items():
            dev = np.nanmedian(np.abs(m[idx][mask] - t) / np.abs(t))
            pull = np.nanmax(np.abs((m[idx][mask] - t) / s[idx][mask]))
            summary_rows.append({"level": lvl, "flavour": FLAV[idx].strip("$\\"),
                                 "method": meth, "median_rel_dev": round(float(dev), 5),
                                 "max_pull": round(float(pull), 3)})

df = pd.DataFrame(summary_rows)
df.to_csv(f"{OUT}/method_comparison_summary.csv", index=False)
print(df.to_string(index=False))