"""Zoomed-in real-vs-generated edep histogram for photons only (pdg=22),
to characterize the KL spike flagged by plot_kl_bars_pl(group_by='pdg')."""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

import giant.analysis as a

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

samples = a.load_predicted_local(FILE, sample_frac=0.15)
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]
print("n photon rows:", mask.sum())
print(
    "real: mean",
    real.mean(),
    "std",
    real.std(),
    "max",
    real.max(),
    "frac==0",
    (real == 0).mean(),
)
print(
    "gen:  mean",
    gen.mean(),
    "std",
    gen.std(),
    "max",
    gen.max(),
    "frac==0",
    (gen == 0).mean(),
)
for q in [0.5, 0.9, 0.99, 0.999]:
    print(f"q={q}: real={np.quantile(real, q):.4f} gen={np.quantile(gen, q):.4f}")

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()
fig.savefig(
    OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-zoom.png", dpi=150, bbox_inches="tight"
)
print("saved photon-edep-zoom")