giant/giant/rollout.py

"""Autoregressive shower rollout driver.

Steps the two-stage GIANT surrogate forward into a full particle shower: each
primary post-step becomes the next pre-step, secondaries are pushed as new tracks,
and the material/layer_id conditioning at every step comes from a `GeometryOracle`
(the surrogate does not predict them).

Tracks are advanced breadth-first: every sweep steps all currently-active tracks
once (in `batch_size` chunks), so many tracks share each model forward pass. A
track terminates on one of the recorded `termination_reason`s in constants.py.

Energy accounting: on every terminal stop except escape, the track's remaining
energy is deposited locally so the shower conserves energy; escaped energy is
treated as detector leakage and not deposited.
"""

from __future__ import annotations

from collections import Counter
from typing import Callable, TypedDict

import numpy as np
import torch

from giant.constants import (
    TERM_ENERGY_CUTOFF,
    TERM_ESCAPED,
    TERM_MAX_STEPS,
    TERM_NATURAL_END,
    TERM_UNKNOWN_PDG,
)
from giant.data.transforms import (
    Normalizer,
    build_cond_features,
    decode_secondaries,
    energy_simplex_decode,
    inv_local_frame_rotation,
    inv_log_transform,
    reconstruct_post_pos,
)
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx

# Record columns produced per step / per terminal marker.
_RECORD_KEYS = [
    "event_id",
    "track_id",
    "parent_id",
    "generation",
    "step_no",
    "pdg",
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "post_x",
    "post_y",
    "post_z",
    "post_E",
    "post_dx",
    "post_dy",
    "post_dz",
    "edep",
    "step_length",
    "material",
    "layer_id",
    "n_sec_pred",
    "termination_reason",
]


def _empty_frontier() -> dict[str, np.ndarray]:
    return {
        "event_id": np.empty(0, dtype=np.int64),
        "track_id": np.empty(0, dtype=np.int64),
        "parent_id": np.empty(0, dtype=np.int64),
        "generation": np.empty(0, dtype=np.int64),
        "step_in_track": np.empty(0, dtype=np.int64),
        "pdg": np.empty(0, dtype=np.int64),
        "pre_pos": np.empty((0, 3), dtype=np.float64),
        "pre_E": np.empty(0, dtype=np.float64),
        "pre_dir": np.empty((0, 3), dtype=np.float64),
    }


def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
    parts = [p for p in parts if len(p["event_id"]) > 0]
    if not parts:
        return _empty_frontier()
    return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]}


# Fixed per-key dtype, so every chunk table has an identical schema — needed
# for `giant rollout --on_chunk` to stream chunks straight into one
# pq.ParquetWriter (which requires matching schemas across writes), and a
# side benefit even in the buffered path since np.concatenate would otherwise
# silently upcast any stray int32/float32 chunk to the majority dtype.
_RECORD_DTYPES: dict[str, type] = {
    "event_id": np.int64,
    "track_id": np.int64,
    "parent_id": np.int64,
    "generation": np.int64,
    "step_no": np.int64,
    "pdg": np.int64,
    "pre_x": np.float64,
    "pre_y": np.float64,
    "pre_z": np.float64,
    "pre_E": np.float64,
    "pre_dx": np.float64,
    "pre_dy": np.float64,
    "pre_dz": np.float64,
    "post_x": np.float64,
    "post_y": np.float64,
    "post_z": np.float64,
    "post_E": np.float64,
    "post_dx": np.float64,
    "post_dy": np.float64,
    "post_dz": np.float64,
    "edep": np.float64,
    "step_length": np.float64,
    "material": object,
    "layer_id": np.int64,
    "n_sec_pred": np.int64,
    "termination_reason": object,
}


class RolloutSummary(TypedDict):
    """`rollout()`'s return shape when streaming to `on_chunk` instead of materializing rows."""

    n_rows: int
    termination_reason_counts: dict[str, int]


class _Recorder:
    """Accumulates per-step rows into column lists, materialised at the end —
    or, when `sink` is given, streams each non-empty chunk to it immediately
    instead, keeping only row-count / termination-reason summaries in memory.

    The streaming path is what lets `giant rollout` write output incrementally
    (see `rollout`'s `on_chunk` parameter): without it, a whole run's steps —
    scaling with `n_events * max_steps * avg_tracks_per_event` — would sit in
    RAM until the very end.
    """

    def __init__(
        self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
    ) -> None:
        self._sink = sink
        self._cols: dict[str, list] | None = (
            None if sink is not None else {k: [] for k in _RECORD_KEYS}
        )
        self.n_rows = 0
        self.termination_reason_counts: Counter[str] = Counter()

    def add(self, **cols) -> None:
        n = len(cols["event_id"])
        if n == 0:
            return
        row = {
            k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
            for k in _RECORD_KEYS
        }
        self.n_rows += n
        reasons = row["termination_reason"]
        nonempty = reasons[reasons != ""]
        if len(nonempty):
            for r, c in zip(*np.unique(nonempty, return_counts=True)):
                self.termination_reason_counts[str(r)] += int(c)

        if self._sink is not None:
            self._sink(row)
        else:
            assert self._cols is not None
            for k in _RECORD_KEYS:
                self._cols[k].append(row[k])

    def to_dict(self) -> dict[str, np.ndarray]:
        assert self._cols is not None, (
            "to_dict() is unavailable when streaming to a sink — use "
            "n_rows/termination_reason_counts instead"
        )
        out = {}
        for k, chunks in self._cols.items():
            if chunks:
                out[k] = np.concatenate(chunks, axis=0)
            else:
                out[k] = np.empty(0, dtype=_RECORD_DTYPES[k])
        return out


def make_seed_frontier(
    event_id: np.ndarray,
    pdg: np.ndarray,
    pre_pos: np.ndarray,
    pre_E: np.ndarray,
    pre_dir: np.ndarray,
) -> tuple[dict[str, np.ndarray], dict[int, int]]:
    """Build the initial frontier from primary entry states.

    Returns (frontier, event_track_count) where the latter tracks the next
    unused track_id per event (each primary gets a fresh id starting from 0).
    """
    event_id = np.asarray(event_id, dtype=np.int64)
    n = len(event_id)
    track_id = np.empty(n, dtype=np.int64)
    counts: dict[int, int] = {}
    for i, ev in enumerate(event_id.tolist()):
        c = counts.get(ev, 0)
        track_id[i] = c
        counts[ev] = c + 1

    dir_ = np.asarray(pre_dir, dtype=np.float64)
    dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)

    frontier = {
        "event_id": event_id,
        "track_id": track_id,
        "parent_id": np.full(n, -1, dtype=np.int64),
        "generation": np.zeros(n, dtype=np.int64),
        "step_in_track": np.zeros(n, dtype=np.int64),
        "pdg": np.asarray(pdg, dtype=np.int64),
        "pre_pos": np.asarray(pre_pos, dtype=np.float64),
        "pre_E": np.asarray(pre_E, dtype=np.float64),
        "pre_dir": dir_,
    }
    return frontier, counts


def _terminal_rows(tr: dict[str, np.ndarray], sel: np.ndarray, reason: str, edep):
    """Assemble terminal-marker record columns for the selected tracks."""
    pos = tr["pre_pos"][sel]
    dir_ = tr["pre_dir"][sel]
    n = int(sel.sum())
    return dict(
        event_id=tr["event_id"][sel],
        track_id=tr["track_id"][sel],
        parent_id=tr["parent_id"][sel],
        generation=tr["generation"][sel],
        step_no=tr["step_in_track"][sel],
        pdg=tr["pdg"][sel],
        pre_x=pos[:, 0],
        pre_y=pos[:, 1],
        pre_z=pos[:, 2],
        pre_E=tr["pre_E"][sel],
        pre_dx=dir_[:, 0],
        pre_dy=dir_[:, 1],
        pre_dz=dir_[:, 2],
        post_x=pos[:, 0],
        post_y=pos[:, 1],
        post_z=pos[:, 2],
        post_E=np.zeros(n),
        post_dx=dir_[:, 0],
        post_dy=dir_[:, 1],
        post_dz=dir_[:, 2],
        edep=np.asarray(edep, dtype=np.float64).reshape(n),
        step_length=np.zeros(n),
        material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
        layer_id=tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
        n_sec_pred=np.zeros(n, dtype=np.int64),
        termination_reason=np.full(n, reason, dtype=object),
    )


@torch.no_grad()
def rollout(
    stage1_model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    oracle,
    seeds: dict[str, np.ndarray],
    cond_norm: Normalizer,
    tgt_norm: Normalizer,
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    *,
    energy_cutoff: float,
    max_steps: int,
    steps: int = 10,
    batch_size: int = 4096,
    device: torch.device | None = None,
    max_tracks_per_event: int | None = None,
    escape_threshold: float | None = None,
    on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
) -> dict[str, np.ndarray] | RolloutSummary:
    """Run showers to completion.

    By default, returns a step-record dict (see _RECORD_KEYS) with the whole
    run's rows materialised in memory.

    If `on_chunk` is given, every non-empty batch of rows is streamed to it as
    soon as it's produced instead — no per-run buffering — and this returns a
    small summary dict instead: `{"n_rows": int, "termination_reason_counts":
    dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs,
    where the full record set would otherwise scale with
    `n_events * max_steps * avg_tracks_per_event`.
    """
    device = device or torch.device("cpu")
    stage1_model.eval()
    sec_decoder.eval()
    if escape_threshold is not None:
        oracle.escape_threshold = float(escape_threshold)

    pdg_map_inv = {v: k for k, v in pdg_map.items()}
    pdg_emb_weight = stage1_model.pdg_embedding_weight()

    frontier, counts = make_seed_frontier(
        seeds["event_id"],
        seeds["pdg"],
        seeds["pre_pos"],
        seeds["pre_E"],
        seeds["pre_dir"],
    )
    rec = _Recorder(sink=on_chunk)

    while len(frontier["event_id"]) > 0:
        next_parts: list[dict[str, np.ndarray]] = []
        n_total = len(frontier["event_id"])
        for start in range(0, n_total, batch_size):
            chunk = {k: v[start : start + batch_size] for k, v in frontier.items()}
            next_parts.append(
                _step_chunk(
                    chunk,
                    stage1_model,
                    sec_decoder,
                    oracle,
                    cond_norm,
                    tgt_norm,
                    pdg_map,
                    mat_map,
                    pdg_map_inv,
                    pdg_emb_weight,
                    rec,
                    counts,
                    energy_cutoff,
                    max_steps,
                    steps,
                    device,
                    max_tracks_per_event,
                )
            )
        frontier = _concat_frontiers(next_parts)

    if on_chunk is not None:
        return {
            "n_rows": rec.n_rows,
            "termination_reason_counts": dict(rec.termination_reason_counts),
        }
    return rec.to_dict()


def _step_chunk(
    tr,
    stage1_model,
    sec_decoder,
    oracle,
    cond_norm,
    tgt_norm,
    pdg_map,
    mat_map,
    pdg_map_inv,
    pdg_emb_weight,
    rec,
    counts,
    energy_cutoff,
    max_steps,
    steps,
    device,
    max_tracks_per_event,
) -> dict[str, np.ndarray]:
    """Advance one chunk of tracks by a single step; return the next frontier."""
    n = len(tr["event_id"])

    # --- Geometry lookup + material/layer conditioning ---
    material, layer_id, escaped = oracle.query(tr["pre_pos"])
    tr = dict(tr)
    tr["_material"] = material
    tr["_layer_id"] = layer_id

    known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool)

    # --- Pre-step termination gates (in priority order; each track picks one) ---
    stop = np.zeros(n, dtype=bool)
    escaped_sel = escaped & ~stop
    rec.add(
        **_terminal_rows(
            tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
        )
    )
    stop |= escaped_sel

    unknown_sel = ~known_pdg & ~stop
    rec.add(
        **_terminal_rows(
            tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
        )
    )
    stop |= unknown_sel

    cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
    rec.add(
        **_terminal_rows(
            tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
        )
    )
    stop |= cutoff_sel

    maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
    rec.add(
        **_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
    )
    stop |= maxstep_sel

    active = ~stop
    if not active.any():
        return _empty_frontier()

    tr = {k: v[active] for k, v in tr.items()}
    material = tr["_material"]
    layer_id = tr["_layer_id"]

    # --- Build conditioning and run the two stages ---
    cond_dict = {
        "pre_pos": tr["pre_pos"],
        "pre_E": tr["pre_E"],
        "pre_dir": tr["pre_dir"],
        "layer_id": layer_id,
        "material": material,
        "pdg": tr["pdg"],
    }
    cond_cont, cond_cat = build_cond_features(cond_dict, pdg_map, mat_map, cond_norm)
    cc = torch.from_numpy(cond_cont).float().to(device)
    ck = torch.from_numpy(cond_cat).long().to(device)

    stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
    raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())

    step_length = inv_log_transform(raw[:, 0])
    edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])

    post_dir_local = raw[:, 3:6].copy()
    post_dir_local /= np.clip(
        np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
    )
    post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)

    travel_dir_local = raw[:, 6:9].copy()
    travel_dir_local /= np.clip(
        np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
    )
    post_pos = reconstruct_post_pos(
        tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
    )

    n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)

    # --- Secondaries ---
    sec_cont, sec_type_emb, _valid = sample_secondaries(
        sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
    )
    sec_pdg_idx = snap_type_to_pdg_idx(sec_type_emb, pdg_emb_weight)
    sec_E, sec_dir_world, sec_pdg_code, sec_valid = decode_secondaries(
        sec_cont.cpu().numpy(),
        sec_pdg_idx.cpu().numpy(),
        n_sec_np,
        e_sec,
        tr["pre_dir"],
        pdg_map_inv,
    )

    edep = edep.astype(np.float64)
    post_E = post_E.astype(np.float64)

    # --- Spawn secondaries (with per-event track cap) ---
    new_tracks, dropped_edep = _spawn_secondaries(
        tr,
        post_pos,
        sec_valid,
        sec_E,
        sec_dir_world,
        sec_pdg_code,
        counts,
        max_tracks_per_event,
    )
    # Energy bookkeeping so each step conserves exactly (edep + carried + post_E
    # == pre_E): `decode_secondaries` already rescales valid slots to sum to
    # exactly `e_sec` whenever n_sec > 0, so `residual` here is ~0 except when
    # n_sec == 0 (no secondary to carry the budget at all — the whole `e_sec`
    # becomes residual). Also deposit the energy of any sub-cap secondaries
    # we dropped for hitting `max_tracks_per_event`.
    sec_E_valid_sum = (sec_E * sec_valid).sum(axis=1)
    residual = np.maximum(e_sec - sec_E_valid_sum, 0.0)
    edep = edep + residual + dropped_edep

    # --- Record the stepped rows; mark natural_end where the primary died ---
    natural = post_E <= 0.0
    reason = np.where(natural, TERM_NATURAL_END, "").astype(object)
    rec.add(
        event_id=tr["event_id"],
        track_id=tr["track_id"],
        parent_id=tr["parent_id"],
        generation=tr["generation"],
        step_no=tr["step_in_track"],
        pdg=tr["pdg"],
        pre_x=tr["pre_pos"][:, 0],
        pre_y=tr["pre_pos"][:, 1],
        pre_z=tr["pre_pos"][:, 2],
        pre_E=tr["pre_E"],
        pre_dx=tr["pre_dir"][:, 0],
        pre_dy=tr["pre_dir"][:, 1],
        pre_dz=tr["pre_dir"][:, 2],
        post_x=post_pos[:, 0],
        post_y=post_pos[:, 1],
        post_z=post_pos[:, 2],
        post_E=post_E,
        post_dx=post_dir_world[:, 0],
        post_dy=post_dir_world[:, 1],
        post_dz=post_dir_world[:, 2],
        edep=edep,
        step_length=step_length,
        material=material,
        layer_id=layer_id,
        n_sec_pred=n_sec_np,
        termination_reason=reason,
    )

    # --- Continue surviving primaries ---
    cont = ~natural
    cont_frontier = {
        "event_id": tr["event_id"][cont],
        "track_id": tr["track_id"][cont],
        "parent_id": tr["parent_id"][cont],
        "generation": tr["generation"][cont],
        "step_in_track": tr["step_in_track"][cont] + 1,
        "pdg": tr["pdg"][cont],
        "pre_pos": post_pos[cont],
        "pre_E": post_E[cont],
        "pre_dir": post_dir_world[cont],
    }
    return _concat_frontiers([cont_frontier, new_tracks])


def _spawn_secondaries(
    tr,
    post_pos,
    sec_valid,
    sec_E,
    sec_dir_world,
    sec_pdg_code,
    counts,
    max_tracks_per_event,
) -> tuple[dict[str, np.ndarray], np.ndarray]:
    """Turn valid secondaries into new tracks; return (frontier, per-parent dropped edep).

    Secondaries are born at their parent's post_pos. When `max_tracks_per_event`
    is set and an event is at its cap, further secondaries are not spawned; their
    energy is returned as `dropped_edep` (indexed by parent row) so it is
    deposited into the parent step instead of vanishing.
    """
    B = len(tr["event_id"])
    dropped_edep = np.zeros(B, dtype=np.float64)

    pr, sl = np.nonzero(sec_valid)  # parent-row idx, slot idx
    if len(pr) == 0:
        return _empty_frontier(), dropped_edep

    # Assign a fresh per-event track_id to each candidate in stable parent order,
    # applying the per-event cap. The candidate count per chunk is small
    # (<= batch_size * K_MAX), so a plain loop is clear and fast enough.
    order = np.lexsort((sl, pr))  # group by parent row, slot ascending
    kept_pr, kept_sl, kept_tid = [], [], []
    for j in order:
        ev = int(tr["event_id"][pr[j]])
        cur = counts.get(ev, 0)
        if max_tracks_per_event is not None and cur >= max_tracks_per_event:
            dropped_edep[pr[j]] += float(sec_E[pr[j], sl[j]])
            continue
        kept_pr.append(pr[j])
        kept_sl.append(sl[j])
        kept_tid.append(cur)
        counts[ev] = cur + 1

    if not kept_pr:
        return _empty_frontier(), dropped_edep

    pr_k = np.array(kept_pr, dtype=np.int64)
    sl_k = np.array(kept_sl, dtype=np.int64)
    tid_k = np.array(kept_tid, dtype=np.int64)

    frontier = {
        "event_id": tr["event_id"][pr_k],
        "track_id": tid_k,
        "parent_id": tr["track_id"][pr_k],
        "generation": tr["generation"][pr_k] + 1,
        "step_in_track": np.zeros(len(pr_k), dtype=np.int64),
        "pdg": sec_pdg_code[pr_k, sl_k].astype(np.int64),
        "pre_pos": post_pos[pr_k],
        "pre_E": sec_E[pr_k, sl_k].astype(np.float64),
        "pre_dir": sec_dir_world[pr_k, sl_k].astype(np.float64),
    }
    return frontier, dropped_edep