"""Feature-by-feature comparison of two SPAtop HDF5 files, compiled into one PDF.

Usage: python compare_features.py <new.h5> <old.h5> <out.pdf> [max_events]
Histograms are unit-normalized (shape comparison); ratio panels show new/old.
Only real objects enter per-object histograms (INPUTS/<coll>/MASK).
"""
import json
import sys

import h5py
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_pdf import PdfPages

NEW, OLD, OUT = sys.argv[1], sys.argv[2], sys.argv[3]
MAXEV = int(sys.argv[4]) if len(sys.argv) > 4 else 50000
LAB = {"new": "collide2v (PUPPI, 200 PU)", "old": "original Delphes (tpm70_wpm30_FB350_fixed_test)"}
COL = {"new": "#3f90da", "old": "#bd1f01"}
TOP, W = 172.5, 80.0


def load(path):
    out = {}
    with h5py.File(path, "r") as f:
        n = min(MAXEV, f["INPUTS/Jets/MASK"].shape[0])
        def visit(name, obj):
            if isinstance(obj, h5py.Dataset):
                out[name] = obj[:n]
        f.visititems(visit)
        attrs = {k: (v if isinstance(v, (int, float, str)) else str(v)) for k, v in f.attrs.items()}
        ntot = f["INPUTS/Jets/MASK"].shape[0]
    return out, n, ntot, attrs


D = {}
D["new"], n_new, ntot_new, attrs_new = load(NEW)
D["old"], n_old, ntot_old, attrs_old = load(OLD)
N = {"new": n_new, "old": n_old}


def p4(pt, eta, phi, m):
    px, py, pz = pt * np.cos(phi), pt * np.sin(phi), pt * np.sinh(eta)
    E = np.sqrt(np.maximum(px ** 2 + py ** 2 + pz ** 2 + np.maximum(m, 0) ** 2, 0))
    return np.stack([px, py, pz, E], -1)


def minv(v):
    return np.sqrt(np.maximum(v[..., 3] ** 2 - v[..., 0] ** 2 - v[..., 1] ** 2 - v[..., 2] ** 2, 0))


def take(arr, idx):
    """arr (N, slots, ...), idx (N,) with -1 for missing -> values (N, ...) and validity."""
    ok = idx >= 0
    v = arr[np.arange(len(idx)), np.clip(idx, 0, arr.shape[1] - 1)]
    return v, ok


def binning(*arrs, nbins=50, lo_q=0.005, hi_q=0.995, integer=False):
    arrs = [np.asarray(x, dtype=float) for x in arrs if x is not None and len(x)]
    allv = np.concatenate(arrs) if arrs else np.array([])
    allv = allv[np.isfinite(allv)]
    if allv.size == 0:
        return np.linspace(0, 1, nbins + 1)
    lo, hi = np.quantile(allv, lo_q), np.quantile(allv, hi_q)
    if integer or np.allclose(allv, np.round(allv)):  # integer-valued: bins centred on integers
        lo_i, hi_i = np.floor(lo), np.ceil(hi)
        width = max(1.0, np.ceil((hi_i - lo_i + 1) / nbins))
        return np.arange(lo_i - 0.5, hi_i + width + 0.5, width)
    if np.unique(allv).size <= 25:
        u = np.unique(allv)
        if u.size <= 60:
            return np.arange(u.min() - 0.5, u.max() + 1.5, 1.0) if np.allclose(u, np.round(u)) else np.linspace(u.min(), u.max() + 1e-6, min(u.size + 1, nbins + 1))
    if hi <= lo:
        hi = lo + 1
    return np.linspace(lo, hi, nbins + 1)


def draw(ax, rax, vals, bins, title, xlabel, logy=False):
    hs = {}
    for k in ("old", "new"):
        x = vals.get(k)
        if x is None or len(x) == 0:
            continue
        x = np.clip(x[np.isfinite(x)], bins[0], bins[-1] - 1e-9)   # under/overflow into edge bins
        h, _ = np.histogram(x, bins=bins)
        hn = h / max(h.sum(), 1) / np.diff(bins)
        err = np.sqrt(h) / max(h.sum(), 1) / np.diff(bins)
        ax.stairs(hn, bins, color=COL[k], label=f"{LAB[k]} (n={len(x)})", lw=1.5)
        hs[k] = (hn, err)
    ax.set_title(title, fontsize=9)
    ax.set_ylabel("normalized", fontsize=8)
    if logy:
        ax.set_yscale("log")
    ax.tick_params(labelsize=7, labelbottom=False)
    ax.legend(fontsize=6, loc="best")
    if rax is not None and "old" in hs and "new" in hs:
        with np.errstate(divide="ignore", invalid="ignore"):
            r = hs["new"][0] / hs["old"][0]
            re = r * np.sqrt((hs["new"][1] / hs["new"][0]) ** 2 + (hs["old"][1] / hs["old"][0]) ** 2)
        c = 0.5 * (bins[1:] + bins[:-1])
        rax.errorbar(c, r, yerr=re, fmt=".", ms=3, color="k", elinewidth=0.6)
        rax.axhline(1, color="gray", lw=0.6)
        rax.set_ylim(0, 2.5)
        rax.set_ylabel("new/old", fontsize=7)
    if rax is not None:
        rax.set_xlabel(xlabel, fontsize=8)
        rax.tick_params(labelsize=7)


class Pages:
    def __init__(self, pdf):
        self.pdf, self.fig, self.k = pdf, None, 0

    def panel(self):
        if self.fig is None or self.k == 4:
            self.flush()
            self.fig = plt.figure(figsize=(11, 8.5))
            self.gs = self.fig.add_gridspec(2, 2, hspace=0.35, wspace=0.25)
            self.k = 0
        sub = self.gs[self.k // 2, self.k % 2].subgridspec(2, 1, height_ratios=[3, 1], hspace=0.05)
        ax = self.fig.add_subplot(sub[0]); rax = self.fig.add_subplot(sub[1], sharex=ax)
        self.k += 1
        return ax, rax

    def flush(self):
        if self.fig is not None:
            self.pdf.savefig(self.fig); plt.close(self.fig); self.fig = None

    def text_page(self, lines):
        self.flush()
        fig = plt.figure(figsize=(11, 8.5)); fig.text(0.05, 0.95, "\n".join(lines), va="top", family="monospace", fontsize=8)
        self.pdf.savefig(fig); plt.close(fig)


def per_object(k, coll, feat):
    d = D[k]
    key = f"INPUTS/{coll}/{feat}"
    if key not in d:
        return None
    m = d[f"INPUTS/{coll}/MASK"].astype(bool)
    return d[key][m].astype(np.float64)


with PdfPages(OUT) as pdf:
    P = Pages(pdf)
    keys_new = {k for k in D["new"] if k.startswith("INPUTS/")}
    keys_old = {k for k in D["old"] if k.startswith("INPUTS/")}
    P.text_page([
        "SPAtop feature comparison: collide2v vs original Delphes",
        "",
        f"new: {NEW}",
        f"     events used {N['new']} of {ntot_new}; git_sha {attrs_new.get('git_sha', '-')}",
        f"old: {OLD}",
        f"     events used {N['old']} of {ntot_old}; attrs {list(attrs_old)[:5]}",
        "",
        "All histograms are normalized to unit area (shape comparison). Under/overflow is put in the edge bins.",
        "Per-object histograms include only real (unpadded) objects. Ratio panels: new / old, with statistical errors.",
        "",
        "Features only in new: " + ", ".join(sorted(k.split('/', 1)[1] for k in keys_new - keys_old)),
        "Features only in old: " + ", ".join(sorted(k.split('/', 1)[1] for k in keys_old - keys_new)),
        "",
        "new settings: " + (attrs_new.get("settings", "-")[:1500]),
    ])

    # 1) every input feature
    for coll in ("Jets", "BoostedJets"):
        feats = sorted({k.split("/")[2] for k in keys_new | keys_old if k.startswith(f"INPUTS/{coll}/") and not k.endswith("/MASK")})
        for feat in feats:
            vals = {k: per_object(k, coll, feat) for k in ("new", "old")}
            vals = {k: v for k, v in vals.items() if v is not None}
            a, b = vals.get("new", np.array([])), vals.get("old", np.array([]))
            logy = feat in ("pt", "fj_pt", "mass", "fj_mass", "deltaRfj", "fj_ehadovereem", "fj_nneutral", "fj_ncharged")
            bins = binning(a, b)
            ax, rax = P.panel()
            draw(ax, rax, vals, bins, f"{coll}/{feat}", feat, logy=logy)

    # 1b) deltaRfj without the 999 "no fatjet within 0.8" sentinel
    vals = {}
    for k in ("new", "old"):
        v = per_object(k, "Jets", "deltaRfj")
        vals[k] = v[v < 900] if v is not None else np.array([])
    ax, rax = P.panel(); draw(ax, rax, vals, np.linspace(0, 0.8, 41), "Jets/deltaRfj for jets with a fatjet within 0.8 (sentinel 999 removed)", "deltaRfj")
    vals = {}
    for k in ("new", "old"):
        v = per_object(k, "Jets", "deltaRfj")
        vals[k] = (v >= 900).astype(float) if v is not None else np.array([])
    ax, rax = P.panel(); draw(ax, rax, vals, np.array([-0.5, 0.5, 1.5]), "Jets/deltaRfj: fraction of jets with no fatjet within 0.8 (value 999)", "1 = sentinel")

    # 2) multiplicities and per-slot pT
    for coll, ptk in (("Jets", "pt"), ("BoostedJets", "fj_pt")):
        vals = {k: D[k][f"INPUTS/{coll}/MASK"].sum(1).astype(float) for k in ("new", "old")}
        ax, rax = P.panel(); draw(ax, rax, vals, binning(*vals.values(), integer=True), f"number of {coll} per event", "n")
        nslots = min(D["new"][f"INPUTS/{coll}/MASK"].shape[1], D["old"][f"INPUTS/{coll}/MASK"].shape[1])
        for s in range(min(nslots, 6 if coll == "Jets" else 3)):
            vals = {k: D[k][f"INPUTS/{coll}/{ptk}"][:, s][D[k][f"INPUTS/{coll}/MASK"][:, s].astype(bool)].astype(float) for k in ("new", "old")}
            ax, rax = P.panel(); draw(ax, rax, vals, binning(*vals.values()), f"{coll} slot {s} {ptk} (pT-ordered)", ptk, logy=True)
    for bt in ("btag", "btag_L", "btag_M", "btag_T"):
        vals = {k: (D[k][f"INPUTS/Jets/{bt}"] & D[k]["INPUTS/Jets/MASK"]).sum(1).astype(float) for k in ("new", "old") if f"INPUTS/Jets/{bt}" in D[k]}
        ax, rax = P.panel(); draw(ax, rax, vals, binning(*[v for v in vals.values()], integer=True), f"number of {bt} jets per event", "n")

    # 3) targets: category fractions, number of valid targets, gen top pT per category
    P.flush()
    tnames = sorted({k.split("/")[1] for k in list(D["new"]) + list(D["old"]) if k.startswith("TARGETS/") and k.endswith("/MASK")})
    fig, ax = plt.subplots(figsize=(11, 5))
    x = np.arange(len(tnames))
    for j, k in enumerate(("old", "new")):
        fr = [D[k][f"TARGETS/{t}/MASK"].mean() if f"TARGETS/{t}/MASK" in D[k] else np.nan for t in tnames]
        ax.bar(x + (j - 0.5) * 0.4, fr, 0.4, color=COL[k], label=LAB[k])
    ax.set_xticks(x); ax.set_xticklabels(tnames, rotation=30); ax.set_ylabel("fraction of events with valid target"); ax.legend(fontsize=8)
    ax.set_title("Target category fractions (TARGETS/<cat>/MASK)")
    pdf.savefig(fig); plt.close(fig)
    vals = {k: np.sum([D[k][f"TARGETS/{t}/MASK"] for t in tnames if f"TARGETS/{t}/MASK" in D[k]], axis=0).astype(float) for k in ("new", "old")}
    ax, rax = P.panel(); draw(ax, rax, vals, binning(*vals.values(), integer=True), "number of valid target labels per event (all categories)", "n")
    for cat in ("FR", "SRqq", "SRbq", "FB"):
        vals = {}
        for k in ("new", "old"):
            v = [D[k][f"TARGETS/{t}/pt"][D[k][f"TARGETS/{t}/MASK"].astype(bool)] for t in tnames if t.startswith(cat + "t") and f"TARGETS/{t}/pt" in D[k]]
            vals[k] = np.concatenate(v).astype(float) if v else np.array([])
        ax, rax = P.panel(); draw(ax, rax, vals, np.linspace(0, 1200, 49), f"gen top pT of valid {cat} targets", "top pT [GeV]")
    vals = {}
    for k in ("new", "old"):
        v = [D[k][f"TARGETS/{t}/pt"] for t in tnames if t.endswith("t1") or t.endswith("t2")]
        allpt = np.concatenate([D[k]["TARGETS/FRt1/pt"], D[k]["TARGETS/FRt2/pt"]]).astype(float)
        vals[k] = allpt[allpt > 0]
    ax, rax = P.panel(); draw(ax, rax, vals, np.linspace(0, 1200, 49), "gen top pT, all hadronic tops in selected events", "top pT [GeV]")

    # 4) reconstructed masses of truth-matched targets (where the mass windows act)
    def jets4(k):
        d = D[k]
        return p4(d["INPUTS/Jets/pt"].astype(float), d["INPUTS/Jets/eta"].astype(float), d["INPUTS/Jets/phi"].astype(float), d["INPUTS/Jets/mass"].astype(float))

    def fjets4(k):
        d = D[k]
        return p4(d["INPUTS/BoostedJets/fj_pt"].astype(float), d["INPUTS/BoostedJets/fj_eta"].astype(float), d["INPUTS/BoostedJets/fj_phi"].astype(float), d["INPUTS/BoostedJets/fj_mass"].astype(float))

    mt, mw, msr, mfb, sdfb, t32fb, sdo, t32o = ({} for _ in range(8))
    for k in ("new", "old"):
        d = D[k]; J = jets4(k); FJ = fjets4(k)
        a_t, a_w, a_sr, a_fb, a_sd, a_t32 = [], [], [], [], [], []
        isfb = np.zeros(d["INPUTS/BoostedJets/MASK"].shape, bool)
        for i in (1, 2):
            if f"TARGETS/FRt{i}/MASK" in d:
                mk = d[f"TARGETS/FRt{i}/MASK"].astype(bool)
                vb, _ = take(J, d[f"TARGETS/FRt{i}/b"]); v1, _ = take(J, d[f"TARGETS/FRt{i}/q1"]); v2, _ = take(J, d[f"TARGETS/FRt{i}/q2"])
                a_t.append(minv(vb + v1 + v2)[mk]); a_w.append(minv(v1 + v2)[mk])
            if f"TARGETS/SRqqt{i}/MASK" in d:
                mk = d[f"TARGETS/SRqqt{i}/MASK"].astype(bool)
                vb, _ = take(J, d[f"TARGETS/SRqqt{i}/b"]); vq, _ = take(FJ, d[f"TARGETS/SRqqt{i}/qq"])
                a_sr.append(minv(vb + vq)[mk])
            if f"TARGETS/FBt{i}/MASK" in d:
                mk = d[f"TARGETS/FBt{i}/MASK"].astype(bool); idx = d[f"TARGETS/FBt{i}/bqq"]
                a_fb.append(d["INPUTS/BoostedJets/fj_mass"][np.arange(len(idx)), np.clip(idx, 0, None)][mk])
                a_sd.append(d["INPUTS/BoostedJets/fj_sdmass"][np.arange(len(idx)), np.clip(idx, 0, None)][mk])
                a_t32.append(d["INPUTS/BoostedJets/fj_tau32"][np.arange(len(idx)), np.clip(idx, 0, None)][mk])
                isfb[np.where(mk)[0], idx[mk]] = True
        cat = lambda L: np.concatenate(L).astype(float) if L else np.array([])
        mt[k], mw[k], msr[k], mfb[k], sdfb[k], t32fb[k] = map(cat, (a_t, a_w, a_sr, a_fb, a_sd, a_t32))
        other = d["INPUTS/BoostedJets/MASK"].astype(bool) & ~isfb
        sdo[k] = d["INPUTS/BoostedJets/fj_sdmass"][other].astype(float); t32o[k] = d["INPUTS/BoostedJets/fj_tau32"][other].astype(float)
    for vals, title, bins in ((mt, "FR targets: m(b q1 q2) of truth-matched jets [window 102.5-242.5]", np.linspace(80, 260, 46)),
                              (mw, "FR targets: m(q1 q2) of truth-matched jets [window 50-110]", np.linspace(40, 120, 41)),
                              (msr, "SRqq targets: m(b-jet + W fatjet) [window 102.5-242.5]", np.linspace(80, 260, 46)),
                              (mfb, "FB targets: fatjet (ungroomed) mass [window 102.5-242.5]", np.linspace(80, 260, 46)),
                              (sdfb, "FB targets: fatjet soft-drop mass", np.linspace(0, 260, 53)),
                              (t32fb, "FB targets: fatjet tau32", np.linspace(0, 1, 41)),
                              (sdo, "non-FB fatjets: soft-drop mass", np.linspace(0, 260, 53)),
                              (t32o, "non-FB fatjets: tau32", np.linspace(0, 1, 41))):
        ax, rax = P.panel(); draw(ax, rax, vals, bins, title, "")
    # 5) summary numbers
    rows = []
    def frac(x): return f"{np.mean(x):.3f}" if np.size(x) else "-"
    for k in ("old", "new"):
        d = D[k]
        fm = d["INPUTS/BoostedJets/MASK"].astype(bool); jm = d["INPUTS/Jets/MASK"].astype(bool)
        nt = np.sum([d[f"TARGETS/{t}/MASK"] for t in tnames if f"TARGETS/{t}/MASK" in d], axis=0)
        t21 = d["INPUTS/BoostedJets/fj_tau21"][fm]; t32 = d["INPUTS/BoostedJets/fj_tau32"][fm]
        rows.append((LAB[k], [
            ("events used", str(N[k])),
            ("mean jets / event", f"{jm.sum(1).mean():.2f}"),
            ("mean fatjets / event", f"{fm.sum(1).mean():.2f}"),
            ("events with 0 fatjets", frac(fm.sum(1) == 0)),
            ("events with 0 valid targets", frac(nt == 0)),
            ("btag rate per jet", frac(d["INPUTS/Jets/btag"][jm])),
            ("fj_Ttag rate per fatjet", frac(d["INPUTS/BoostedJets/fj_Ttag"][fm])),
            ("fj_Wtag rate per fatjet", frac(d["INPUTS/BoostedJets/fj_Wtag"][fm])),
            ("jets with deltaRfj == 999", frac(d["INPUTS/Jets/deltaRfj"][jm] >= 900)),
            ("deltaRfj integer-valued?", str(bool(np.allclose(d["INPUTS/Jets/deltaRfj"][jm], np.round(d["INPUTS/Jets/deltaRfj"][jm]))))),
            ("fatjets with tau21 >= 0.95", frac(t21 >= 0.95)),
            ("fatjets with tau32 >= 0.95", frac(t32 >= 0.95)),
            ("fatjets with tau21 or tau32 == -1", frac((t21 == -1) | (t32 == -1))),
            ("median FR m(bqq) [GeV]", f"{np.median(mt[k]):.1f}" if mt[k].size else "-"),
            ("median FR m(qq) [GeV]", f"{np.median(mw[k]):.1f}" if mw[k].size else "-"),
            ("median FB fatjet m_SD [GeV]", f"{np.median(sdfb[k]):.1f}" if sdfb[k].size else "-"),
        ] + [(f"{t} MASK fraction", frac(d[f"TARGETS/{t}/MASK"])) for t in tnames if f"TARGETS/{t}/MASK" in d]))
    lines = ["Summary numbers", ""] + [f"{'quantity':40s} {'old':>14s} {'new':>14s}"]
    for (name, _), (_, v_new) in zip(rows[0][1], rows[1][1]):
        lines.append(f"{name:40s} {dict(rows[0][1])[name]:>14s} {v_new:>14s}")
    P.text_page(lines)
    P.flush()
print("wrote", OUT)