"""Truth-only histogram of deposited photon (pdg=22) energy in eV, grouped by
the Geant4 physics process that ended the step. Step-type histograms, one per
process, in both linear and log energy space.

Data is the raw miniCaloSim steps parquet (truth), loaded lazily with polars —
only the photon rows and the (process, edep) columns are materialized.
"""

from pathlib import Path

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

FILE = "/home/lars/Programming/minicalo-data-exploration/pbwo4_10000events_hits.parquet"
OUT = Path("/home/lars/knowledge-base/meta/attachments")

MEV_TO_EV = 1e6

# Lazy load: keep only photon steps and the two columns we need.
df = (
    pl.scan_parquet(FILE)
    .filter(pl.col("pdg") == 22)
    .select(
        pl.col("process"),
        (pl.col("edep") * MEV_TO_EV).alias("edep_ev"),
    )
    .collect()
)
print("n photon rows:", df.height)

# Process order by abundance, so the legend is stable and the busiest on top.
processes = (
    df.group_by("process")
    .len()
    .sort("len", descending=True)
    .get_column("process")
    .to_list()
)
series = {
    p: df.filter(pl.col("process") == p).get_column("edep_ev").to_numpy()
    for p in processes
}

all_edep = df.get_column("edep_ev").to_numpy()
lin_bins = np.linspace(0, np.quantile(all_edep, 0.999), 80)
pos = all_edep[all_edep > 0]
log_bins = np.geomspace(max(pos.min(), 1e-3), pos.max(), 80)


def plot(bins, xscale: str, fname: str) -> None:
    fig, ax = plt.subplots(figsize=(7, 4.5))
    lo, hi = bins[0], bins[-1]
    for p in processes:
        v = series[p]
        # Skip processes with nothing inside the bin range (e.g. all-zero edep
        # processes on the log axis), which would add phantom legend entries.
        if not np.any((v >= lo) & (v <= hi)):
            continue
        ax.hist(v, bins=bins, histtype="step", density=True, label=f"{p} (n={v.size})")
    ax.set_xscale(xscale)
    ax.set_yscale("log")
    ax.set_xlabel("deposited energy [eV] (photons, pdg=22)")
    ax.set_ylabel("density")
    ax.legend(fontsize=8, title="process")
    fig.tight_layout()
    fig.savefig(OUT / fname, dpi=150, bbox_inches="tight")
    print("saved", fname)


plot(lin_bins, "linear", "giant-photon-edep-by-process-ev-lin.png")
plot(log_bins, "log", "giant-photon-edep-by-process-ev-log.png")