"""One-off export of the energy-budget-violation plot for the
energy-conservation PoC checkpoint (ALR simplex output space, retrained on
the patched-Geant4 regenerated dataset). Not part of the package; run
manually.

Unlike `plot_total_energy` (per-step conservation only, no reference to the
fixed primary/incident energy), this adds an explicit comparison against
`primary_E` (== max(pre_E) per event, since this PoC dataset uses a single
fixed incident energy) to show event-level conservation violations that the
per-step ALR simplex constraint does not prevent.
"""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import polars as pl

from giant.analysis import _edep_pl, _hist_edges

FILE = "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet"
OUT = Path("/home/lars/knowledge-base/meta/attachments")
PREFIX = "giant-energy-conservation-poc"

print("=== per-event primary energy + total edep ===")
per_event = (
    pl.scan_parquet(FILE)
    .group_by("event_id")
    .agg(
        pl.col("pre_E").max().alias("primary_E"),
        _edep_pl("true").sum().alias("real_total_edep"),
        _edep_pl("pred").sum().alias("gen_total_edep"),
    )
    .collect(engine="streaming")
)
print("n events:", per_event.height)

primary_E = per_event["primary_E"].to_numpy()
real_tot = per_event["real_total_edep"].to_numpy()
gen_tot = per_event["gen_total_edep"].to_numpy()

assert np.unique(primary_E).size == 1, "expected a single fixed incident energy"
E0 = float(primary_E[0])
print(f"fixed incident energy E0 = {E0} MeV")

print()
print(
    f"real:  mean={real_tot.mean():.3f} std={real_tot.std():.3f} "
    f"sigma/mu={real_tot.std() / real_tot.mean():.4f} max={real_tot.max():.3f} "
    f"frac>E0={np.mean(real_tot > E0):.4f}"
)
print(
    f"gen:   mean={gen_tot.mean():.3f} std={gen_tot.std():.3f} "
    f"sigma/mu={gen_tot.std() / gen_tot.mean():.4f} max={gen_tot.max():.3f} "
    f"frac>E0={np.mean(gen_tot > E0):.4f}"
)
print(
    f"gen/E0 ratio: mean={np.mean(gen_tot / E0):.4f} max={np.max(gen_tot / E0):.4f} "
    f"p99={np.quantile(gen_tot / E0, 0.99):.4f}"
)

print("=== plot: total edep per event, marked against incident energy ===")
fig, ax = plt.subplots(figsize=(6, 4))
edges = _hist_edges(real_tot, gen_tot, bins=50).tolist()
ax.hist(
    real_tot,
    bins=edges,
    density=True,
    histtype="step",
    label=f"real (σ/μ={real_tot.std() / real_tot.mean():.3f})",
)
ax.hist(
    gen_tot,
    bins=edges,
    density=True,
    histtype="step",
    label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, "
    f"{np.mean(gen_tot > E0):.1%} > E0)",
)
ax.axvline(
    E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV"
)
ax.set_yscale("log")
ax.set_xlabel("total deposited energy per event [MeV]")
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(
    OUT / f"{PREFIX}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight"
)

print("=== plot: total edep / incident energy ratio ===")
fig, ax = plt.subplots(figsize=(6, 4))
ratio_real = real_tot / E0
ratio_gen = gen_tot / E0
edges = _hist_edges(ratio_real, ratio_gen, bins=60).tolist()
ax.hist(ratio_real, bins=edges, density=True, histtype="step", label="real")
ax.hist(ratio_gen, bins=edges, density=True, histtype="step", label="generated")
ax.axvline(1.0, color="k", linestyle="--", linewidth=1, label="conservation limit (=1)")
ax.set_yscale("log")
ax.set_xlabel("total deposited energy / incident energy, per event")
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(OUT / f"{PREFIX}-event-energy-ratio.png", dpi=150, bbox_inches="tight")

print("DONE")