"""Export shower-level plots from a `giant rollout` steps parquet.

Not part of the package; run manually. Optionally overlays the real showers
seeded from the same events (a `giant predict --coord local` file) by passing a
reference path. Usage::

    python analysis/export_rollout_observables.py ROLLOUT.parquet [REFERENCE_local.parquet]
"""

import sys
from pathlib import Path

import giant.analysis as a

rollout_file = sys.argv[1] if len(sys.argv) > 1 else "rollout.parquet"
reference_file = (
    sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] not in ("", "-") else None
)
OUT = Path(sys.argv[3]) if len(sys.argv) > 3 else Path(".")

print(f"=== computing rollout observables: {rollout_file} ===")
obs = a.compute_rollout_observables(rollout_file)
tbl = obs.event_table
print(f"n events: {len(tbl)}")
print(
    f"total_edep/event: mean={tbl['total_edep'].mean():.4g} MeV  "
    f"leaked_E/event: mean={tbl['leaked_E'].mean():.4g} MeV  "
    f"n_tracks/event: mean={tbl['n_tracks'].mean():.1f}  "
    f"n_steps/event: mean={tbl['n_steps'].mean():.1f}"
)

reference = None
if reference_file is not None:
    print(f"=== computing real reference: {reference_file} ===")
    reference = a.compute_event_observables_pl(reference_file)

print("=== plots ===")
for name, fn in [
    ("longitudinal", a.plot_rollout_longitudinal),
    ("transverse", a.plot_rollout_transverse),
    ("total-energy", a.plot_rollout_total_energy),
]:
    fig = fn(obs, reference=reference)
    path = OUT / f"rollout-{name}.png"
    fig.savefig(path, dpi=150, bbox_inches="tight")
    print(f"wrote {path}")

print("DONE")