"""One-off export of validation plots for checkpoints/scan/h1024_n8_d0.1_lr0.0003/best.pt
into the knowledge-base attachments folder. Not part of the package; run manually."""

from pathlib import Path

import giant.analysis as a

FILE = "/home/lars/Programming/giant/pbwo4_10k_9_predicted_local.parquet"
OUT = Path("/home/lars/knowledge-base/meta/attachments")
PREFIX = "giant-h1024n8d0.1lr3e-4"

PDG_NAMES = {
    11: "e-",
    -11: "e+",
    22: "gamma",
    2112: "n",
    2212: "p",
}


def pdg_label(code: int) -> str:
    if code in PDG_NAMES:
        return PDG_NAMES[code]
    if code > 1000000000:
        return f"ion{code}"
    return str(code)


print("=== KL bar plots (lazy, full dataset) ===")
for grouping in [None, "energy", "pdg", "material"]:
    fig = a.plot_kl_bars_pl(FILE, group_by=grouping)
    name = f"{PREFIX}-kl-bars-{grouping or 'all'}.png"
    fig.savefig(OUT / name, dpi=150, bbox_inches="tight")
    print("saved", name)

print("=== loading sampled SampleCollection ===")
samples = a.load_predicted_local(FILE, sample_frac=0.15)
print("n rows:", len(samples.gen_raw))

print("=== marginals ===")
fig = a.plot_marginals(samples)
fig.savefig(OUT / f"{PREFIX}-marginals-all.png", dpi=150, bbox_inches="tight")

fig = a.plot_marginals(samples, group_by="energy")
fig.savefig(OUT / f"{PREFIX}-marginals-energy.png", dpi=150, bbox_inches="tight")

fig = a.plot_marginals(samples, group_by="pdg")
fig.savefig(OUT / f"{PREFIX}-marginals-pdg.png", dpi=150, bbox_inches="tight")
print("saved marginals")

print("=== correlation matrices ===")
fig = a.plot_correlation_matrices(samples)
fig.savefig(OUT / f"{PREFIX}-correlation.png", dpi=150, bbox_inches="tight")

print("=== pairwise ===")
fig = a.plot_pairwise(samples, n_sample=5000)
fig.savefig(OUT / f"{PREFIX}-pairwise.png", dpi=150, bbox_inches="tight")

print("=== direction alignment ===")
fig = a.plot_direction_alignment(samples)
fig.savefig(OUT / f"{PREFIX}-direction-alignment.png", dpi=150, bbox_inches="tight")

print("=== constraint violations ===")
fig = a.plot_constraint_violations(samples)
fig.savefig(OUT / f"{PREFIX}-constraint-violations.png", dpi=150, bbox_inches="tight")

print("=== marginal_table aggregate ===")
agg = a.marginal_table(samples)
print(agg.to_string())

print("=== marginal_table by pdg (top rows incl. photon) ===")
by_pdg = a.marginal_table(samples, group_by="pdg")
by_pdg["particle"] = by_pdg["group"].astype(str)
print(by_pdg.to_string())

print("=== photon-only rows ===")
photon_rows = by_pdg[by_pdg["group"].astype(str) == "pdg=22"]
print(photon_rows.to_string())

print("DONE")