"""One-off export of presentation-specific plots for the 2026-06-24 ETP group
update, written directly into the presentation's images/ folder. Not part of
the package; run manually.

All plots are saved as PDF (vector) except the pairwise scatter plot, which
stays PNG/raster in the .tex (scatter plots with thousands of points blow up
as vector files and gain nothing from being scalable).

1. Aggregate marginals, split into a 3x3 grid (instead of plot_marginals'
   single wide row of 9 columns) for a more manageable slide aspect ratio.
2. post_dir/travel_dir unit-norm histograms only (subset of
   plot_constraint_violations).
3. KL bar plot stratified by pdg.
4. Zoomed-in photon edep histogram + CDF.
5. Event-level profiles/histograms/pdg-share plots.
"""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

import giant.analysis as a
from giant.analysis import RAW_TARGET_NAMES, _hist_edges

FILE = "/home/lars/Programming/giant/pbwo4_10k_9_predicted_local.parquet"
OUT = Path(
    "/home/lars/Programming/thesis-presentations/presentations/2026-06-24-group-update/images"
)
PREFIX = "giant-h1024n8d0.1lr3e-4"


def savefig(fig, name: str) -> None:
    fig.savefig(OUT / f"{PREFIX}-{name}.pdf", 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, split 3x3 grid ===")
n_rows, n_cols = 3, 3
fig, axes = plt.subplots(n_rows, n_cols, figsize=(3.6 * n_cols, 2.8 * n_rows))
for idx, name in enumerate(RAW_TARGET_NAMES):
    row, col = divmod(idx, n_cols)
    ax = axes[row][col]
    real, gen = samples.real_raw[:, idx], samples.gen_raw[:, idx]
    edges = _hist_edges(real, gen, bins=50)
    ax.hist(real, bins=edges, density=True, histtype="step", label="real")
    ax.hist(gen, bins=edges, density=True, histtype="step", label="generated")
    ax.set_yscale("log")
    ax.set_title(name, fontsize=9)
    if idx == 0:
        ax.legend(fontsize=7)
fig.tight_layout()
savefig(fig, "marginals-grid")

print("=== post_dir / travel_dir unit-norm histograms ===")
gen = samples.gen_raw
post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)

fig, axes = plt.subplots(1, 2, figsize=(8, 3.5))
for ax, norm, title in [
    (axes[0], post_norm, r"$\|\mathrm{post\_dir}\|$"),
    (axes[1], travel_norm, r"$\|\mathrm{travel\_dir}\|$"),
]:
    ax.hist(norm, bins=_hist_edges(norm, bins=50), histtype="step")
    ax.set_yscale("log")
    ax.axvline(1.0, color="k", linestyle="--", linewidth=1)
    ax.set_title(title)
fig.tight_layout()
savefig(fig, "direction-norms")

print("=== KL bars by pdg ===")
fig = a.plot_kl_bars_pl(FILE, group_by="pdg")
savefig(fig, "kl-bars-pdg")

print("=== photon edep zoom ===")
mask = samples.pdg == 22
edep_idx = a.RAW_TARGET_NAMES.index("edep")
real = samples.real_raw[mask, edep_idx]
gen = samples.gen_raw[mask, edep_idx]

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
bins = np.linspace(0, np.quantile(real, 0.999), 80)
axes[0].hist(real, bins=bins, alpha=0.6, label="real", density=True)
axes[0].hist(gen, bins=bins, alpha=0.6, label="gen", density=True)
axes[0].set_yscale("log")
axes[0].set_xlabel("edep (photons, pdg=22)")
axes[0].legend()
axes[1].hist(
    real,
    bins=bins,
    alpha=0.6,
    label="real",
    density=True,
    cumulative=True,
    histtype="step",
)
axes[1].hist(
    gen,
    bins=bins,
    alpha=0.6,
    label="gen",
    density=True,
    cumulative=True,
    histtype="step",
)
axes[1].set_xlabel("edep (photons, pdg=22) - CDF")
axes[1].legend()
fig.tight_layout()
savefig(fig, "photon-edep-zoom")

print("=== event observables ===")
obs = a.compute_event_observables_pl(FILE)

fig = a.plot_total_energy(obs)
savefig(fig, "event-total-energy")

fig = a.plot_total_length(obs)
savefig(fig, "event-total-length")

fig = a.plot_mean_energy_per_step(obs)
savefig(fig, "event-mean-energy-per-step")

fig = a.plot_mean_length_per_step(obs)
savefig(fig, "event-mean-length-per-step")

fig = a.plot_longitudinal_profile(obs)
savefig(fig, "event-longitudinal-profile")

fig = a.plot_transverse_profile(obs)
savefig(fig, "event-transverse-profile")

fig = a.plot_shower_max_depth(obs)
savefig(fig, "event-shower-max-depth")

print("=== pdg length share ===")
pdg_table = a.pdg_contribution_table_pl(FILE)
fig = a.plot_pdg_length_share(pdg_table)
savefig(fig, "pdg-length-share")

print("DONE")