giant/giant/cli.py

from collections import Counter
from datetime import date, datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Optional
import uuid as uuid_mod

import numpy as np
import yaml
import torch
import typer
from typing_extensions import Annotated

import pyarrow as pa
import pyarrow.parquet as pq
from tqdm import tqdm

from giant import config as gconfig
from giant.constants import (
    LOCAL_TARGET_NAMES,
    PREDICT_COORD_METADATA_KEY,
    PREDICT_SCHEMA_VERSION,
    PREDICT_SCHEMA_VERSION_KEY,
    ROLLOUT_COORD_VALUE,
)
from giant.data.loader import (
    find_parquet_files,
    iter_file_chunks,
    iter_cond_chunks,
)
from giant.data.transforms import (
    build_features,
    build_cond_features,
    decode_secondaries,
    energy_simplex_decode,
    inv_local_frame_rotation,
    inv_log_transform,
    reconstruct_post_pos,
    Normalizer,
)
from giant.geometry import GeometryOracle
from giant.model.network import build_models
from giant.pipeline import run_train_job
from giant.rollout import rollout as run_rollout
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx

app = typer.Typer(no_args_is_help=True)


def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]:
    """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.

    Routed models spend their FLOPs in the (smaller) expert trunks, not the
    monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert
    dims instead when routing is enabled.
    """
    router_cfg = model_cfg.get("router")
    if router_cfg and router_cfg.get("enabled"):
        return (
            model_cfg.get("expert_hidden_dim", 128),
            model_cfg.get("expert_n_blocks", 3),
        )
    return model_cfg["hidden_dim"], model_cfg["n_blocks"]


def _coerce_scalar(value: str) -> object:
    """Best-effort str -> bool/int/float, else leave as str.

    CLI flag values always arrive as strings; router kwargs like
    `n_experts` (int) or `temperature` (float) need to come out typed the
    same way a TOML file's native types would, since they're merged into
    the same `model.router` dict as file-sourced config.
    """
    if value.lower() in ("true", "false"):
        return value.lower() == "true"
    try:
        return int(value)
    except ValueError:
        pass
    try:
        return float(value)
    except ValueError:
        pass
    return value


def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
    """Parse repeated `--router-axis "type:key=val,key=val"` flags into
    `axis{i}_{field}` flat keys (see `_parse_composed_axes` in
    giant.model.network), indexed by flag order — the Nth `--router-axis`
    becomes axis N.
    """
    out: dict[str, object] = {}
    for i, spec in enumerate(specs):
        axis_type, _, rest = spec.partition(":")
        out[f"axis{i}_type"] = axis_type
        for pair in filter(None, rest.split(",")):
            key, _, val = pair.partition("=")
            out[f"axis{i}_{key}"] = _coerce_scalar(val)
    return out


_CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")


def _resolve_prediction_output(data: Path, out: Path | None) -> tuple[Path, Path, str]:
    """Return (out_path, resolved_dataset_path, pred_uuid).

    When *out* is None the output path is derived from *data*:
    - under /ceph/ → fixed central store with a UUID filename
    - elsewhere   → sibling of *data* with a UUID filename
    """
    dataset_path = data.resolve()
    pred_uuid = str(uuid_mod.uuid4())
    if out is None:
        if str(dataset_path).startswith("/ceph/"):
            out = _CEPH_PREDICTIONS / f"{pred_uuid}.parquet"
        else:
            out = data.parent / f"{pred_uuid}.parquet"
    return out, dataset_path, pred_uuid


def _write_prediction_ref(
    checkpoint: Path,
    pred_uuid: str,
    out: Path,
    dataset_path: Path,
    comment: str | None = None,
) -> Path:
    """Write a YAML sidecar in the checkpoint directory and return its path."""
    ref = {
        "prediction_id": pred_uuid,
        "output": str(out),
        "dataset": str(dataset_path),
        "checkpoint": str(checkpoint.resolve()),
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    if comment is not None:
        ref["comment"] = comment
    ref_path = checkpoint.parent / f"{pred_uuid}.yaml"
    ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
    return ref_path


@app.callback()
def _main() -> None:
    """GIANT — Geant4 step-function surrogate."""


class Mode(str, Enum):
    flow = "flow"
    ddpm = "ddpm"


class Coord(str, Enum):
    global_ = "global"
    local = "local"


@app.command()
def train(
    data: Annotated[
        Path, typer.Argument(help="Parquet file or directory of parquet files")
    ],
    config: Annotated[
        Optional[Path],
        typer.Option(
            "--config", "-c", help="TOML config file (overridden by explicit flags)"
        ),
    ] = None,
    mode: Annotated[
        Optional[Mode],
        typer.Option("--mode", "-m", help="Generative model: flow matching or DDPM"),
    ] = None,
    epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
    batch_size: Annotated[
        Optional[str],
        typer.Option(
            "--batch-size",
            "-b",
            help="Integer, or 'auto' to estimate from free GPU memory "
            "(cuda devices only)",
        ),
    ] = None,
    lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
    warmup_epochs: Annotated[
        Optional[int], typer.Option("--warmup-epochs", "-w")
    ] = None,
    hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
    n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
    emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
    dropout: Annotated[
        Optional[float],
        typer.Option(
            "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
        ),
    ] = None,
    router: Annotated[
        Optional[bool],
        typer.Option(
            "--router/--no-router",
            help="Route both stages through a mixture of small experts "
            "instead of one monolithic trunk (see model.router in config.toml)",
        ),
    ] = None,
    router_type: Annotated[
        Optional[str],
        typer.Option(
            "--router-type", help="Router implementation name (see ROUTER_REGISTRY)"
        ),
    ] = None,
    n_experts: Annotated[
        Optional[int], typer.Option("--n-experts", help="Number of routed experts")
    ] = None,
    router_axis: Annotated[
        Optional[list[str]],
        typer.Option(
            "--router-axis",
            help="Composed-router axis spec 'type:key=val,key=val' (repeatable; "
            "Nth flag = axis N). Use with --router-type composed instead of "
            "--n-experts, e.g. --router-axis 'energy:n_experts=4' "
            "--router-axis 'pdg:n_experts=3,emb_dim=8'",
        ),
    ] = None,
    val_fraction: Annotated[
        Optional[float], typer.Option("--val-fraction", "-f")
    ] = None,
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", "-s", help="Random seed for reproducibility"),
    ] = None,
    validate_every: Annotated[
        Optional[int],
        typer.Option(
            "--validate-every",
            "-v",
            help="Run marginal+KL validation every N epochs (0 disables)",
        ),
    ] = None,
    validate_steps: Annotated[
        Optional[int],
        typer.Option(
            "--validate-steps",
            "-t",
            help="Flow matching ODE steps used during marginal validation "
            "(ignored in ddpm mode, which always runs the full schedule)",
        ),
    ] = None,
    shuffle_buffer: Annotated[
        int,
        typer.Option(
            "--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"
        ),
    ] = 65536,
    out: Annotated[
        Optional[Path],
        typer.Option(
            "--out", "-o", help="Checkpoint dir (default: auto from hyperparams)"
        ),
    ] = None,
    device: Annotated[
        Optional[str],
        typer.Option("--device", "-D", help="cpu | cuda | mps (default: auto)"),
    ] = None,
    num_workers: Annotated[Optional[int], typer.Option("--num-workers", "-j")] = None,
    resume: Annotated[
        Optional[Path],
        typer.Option("--resume", "-r", help="Checkpoint .pt to resume training from"),
    ] = None,
) -> None:
    """Train the GIANT surrogate model."""
    batch_size_auto = False
    batch_size_value: Optional[int] = None
    if batch_size is not None:
        if batch_size.strip().lower() == "auto":
            batch_size_auto = True
        else:
            try:
                batch_size_value = int(batch_size)
            except ValueError:
                typer.echo(
                    f"error: --batch-size must be an integer or 'auto', "
                    f"got {batch_size!r}",
                    err=True,
                )
                raise typer.Exit(1)

    cli_train = {
        k: v
        for k, v in {
            "mode": mode.value if mode is not None else None,
            "epochs": epochs,
            "batch_size": batch_size_value,
            "lr": lr,
            "warmup_epochs": warmup_epochs,
            "val_fraction": val_fraction,
            "num_workers": num_workers,
            "seed": seed,
            "validate_every": validate_every,
            "validate_steps": validate_steps,
        }.items()
        if v is not None
    }
    cli_model: dict[str, object] = {
        k: v
        for k, v in {
            "hidden_dim": hidden_dim,
            "n_blocks": n_blocks,
            "emb_dim": emb_dim,
            "dropout": dropout,
        }.items()
        if v is not None
    }
    cli_router: dict[str, object] = {
        k: v
        for k, v in {
            "enabled": router,
            "type": router_type,
            "n_experts": n_experts,
        }.items()
        if v is not None
    }
    if router_axis:
        cli_router.update(_parse_router_axis_flags(router_axis))
    if cli_router:
        cli_model["router"] = cli_router
    cfg = gconfig.merge_cli_overrides(
        gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
    )
    t, m = cfg["train"], cfg["model"]

    _device = torch.device(device) if device else gconfig.auto_device()

    if batch_size_auto:
        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m)
        try:
            t["batch_size"] = gconfig.estimate_batch_size(
                est_hidden_dim, est_n_blocks, _device
            )
        except ValueError as exc:
            typer.echo(f"error: {exc}", err=True)
            raise typer.Exit(1)
        typer.echo(
            f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
        )

    out_dir = out or Path(
        f"checkpoints/{date.today().strftime('%Y%m%d')}"
        f"_{t['mode']}"
        f"_h{m['hidden_dim']}"
        f"_b{m['n_blocks']}"
        f"_e{m['emb_dim']}"
        f"_lr{t['lr']}"
        f"_bs{t['batch_size']}"
    )

    typer.echo(f"device:  {_device}")
    typer.echo(f"out_dir: {out_dir}")

    run_train_job(
        data=data,
        cfg=cfg,
        out_dir=out_dir,
        device=_device,
        shuffle_buffer=shuffle_buffer,
        num_workers=t["num_workers"],
        resume=resume,
        echo=typer.echo,
    )


@app.command()
def predict(
    data: Annotated[
        Path, typer.Argument(help="Parquet file or directory of parquet files")
    ],
    checkpoint: Annotated[
        Path,
        typer.Option(
            "--checkpoint",
            "-c",
            help="Path to checkpoint .pt file (best.pt or last.pt)",
        ),
    ],
    coord: Annotated[
        Coord,
        typer.Option(
            "--coord",
            "-C",
            help="global: full physical units, world frame (default). "
            "local: raw 9D model output (denormalised only, local frame, "
            "log-scaled scalars) alongside the matching ground-truth target "
            "for the same input file — requires post-step columns.",
        ),
    ] = Coord.global_,
    out: Annotated[
        Optional[Path],
        typer.Option(
            "--out",
            "-o",
            help="Output parquet path (default: <data>_predicted[_local].parquet)",
        ),
    ] = None,
    batch_size: Annotated[
        str,
        typer.Option(
            "--batch-size",
            "-b",
            help="Inference batch size, or 'auto' to estimate from free GPU "
            "memory (cuda devices only)",
        ),
    ] = "4096",
    steps: Annotated[
        int, typer.Option("--steps", "-s", help="Flow matching ODE steps")
    ] = 10,
    device: Annotated[
        Optional[str],
        typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"),
    ] = None,
    comment: Annotated[
        Optional[str],
        typer.Option(
            "--comment",
            "-m",
            help="Free-text note recorded in the prediction's YAML sidecar",
        ),
    ] = None,
) -> None:
    """Run trained model on a parquet file and save predictions."""
    batch_size_auto = False
    batch_size_value: Optional[int] = None
    if batch_size.strip().lower() == "auto":
        batch_size_auto = True
    else:
        try:
            batch_size_value = int(batch_size)
        except ValueError:
            typer.echo(
                f"error: --batch-size must be an integer or 'auto', got {batch_size!r}",
                err=True,
            )
            raise typer.Exit(1)

    _device = torch.device(device) if device else gconfig.auto_device()
    typer.echo(f"device: {_device}")

    # --- Load checkpoint ---
    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
    if "model_config" not in ckpt:
        typer.echo(
            "error: checkpoint has no model_config — retrain with the current code",
            err=True,
        )
        raise typer.Exit(1)

    if "sec_decoder" not in ckpt:
        typer.echo(
            "error: checkpoint has no sec_decoder — retrain with the current code",
            err=True,
        )
        raise typer.Exit(1)

    model_cfg = ckpt["model_config"]

    if batch_size_auto:
        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg)
        try:
            batch_size_value = gconfig.estimate_batch_size(
                est_hidden_dim,
                est_n_blocks,
                _device,
                training=False,
            )
        except ValueError as exc:
            typer.echo(f"error: {exc}", err=True)
            raise typer.Exit(1)
        typer.echo(
            f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
        )

    assert batch_size_value is not None
    bs = batch_size_value
    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
    pdg_map_inv = {v: k for k, v in pdg_map.items()}
    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])

    model, sec_decoder = build_models(model_cfg)
    model.load_state_dict(ckpt["model"])
    model.to(_device).eval()

    sec_decoder.load_state_dict(ckpt["sec_decoder"])
    sec_decoder.to(_device).eval()

    typer.echo(f"loaded checkpoint: {checkpoint}")
    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)

    # --- Output path ---
    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
    out.parent.mkdir(parents=True, exist_ok=True)
    typer.echo(f"output: {out}")

    # --- Stream input, generate predictions, write output ---
    files = find_parquet_files(data)
    typer.echo(f"found {len(files)} parquet file(s)")

    writer: pq.ParquetWriter | None = None
    total = 0
    skipped = 0
    unknown_pdg_counts: Counter[int] = Counter()
    total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files)
    chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks

    def _concat(
        a: dict[str, np.ndarray], b: dict[str, np.ndarray]
    ) -> dict[str, np.ndarray]:
        return {k: np.concatenate([a[k], b[k]], axis=0) for k in a}

    def _process(piece: dict[str, np.ndarray]) -> None:
        nonlocal writer, total

        if coord == Coord.local:
            cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(
                piece, pdg_map, mat_map
            )
            cond_cont = cond_norm.transform(cond_cont)
        else:
            cond_cont, cond_cat = build_cond_features(
                piece, 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(model, cc, ck, steps=steps)

        if coord == Coord.global_:
            sec_cont, sec_type_emb, _sec_valid_pred = sample_secondaries(
                sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
            )
            sec_pdg_idx = snap_type_to_pdg_idx(
                sec_type_emb, model.pdg_embedding_weight()
            )
            sec_cont_np = sec_cont.cpu().numpy()
            sec_pdg_idx_np = sec_pdg_idx.cpu().numpy()

        n_sec_pred_np = n_sec_pred.cpu().numpy()
        pred = stage1_norm.cpu().numpy()  # normalised

        # Inverse-normalise → local frame, log-scaled scalars
        raw = tgt_norm.inverse_transform(pred)

        if coord == Coord.local:
            table = pa.table(
                {
                    "event_id": piece["event_id"],
                    "pdg": piece["pdg"],
                    "pre_x": piece["pre_pos"][:, 0],
                    "pre_y": piece["pre_pos"][:, 1],
                    "pre_z": piece["pre_pos"][:, 2],
                    "pre_E": piece["pre_E"],
                    "pre_dx": piece["pre_dir"][:, 0],
                    "pre_dy": piece["pre_dir"][:, 1],
                    "pre_dz": piece["pre_dir"][:, 2],
                    "material": piece["material"],
                    "layer_id": piece["layer_id"],
                    "n_sec": piece["n_sec"],
                    **{
                        f"pred_{name}": raw[:, j]
                        for j, name in enumerate(LOCAL_TARGET_NAMES)
                    },
                    **{
                        f"true_{name}": target_raw[:, j]
                        for j, name in enumerate(LOCAL_TARGET_NAMES)
                    },
                }
            )
        else:
            step_length = inv_log_transform(raw[:, 0])
            # Columns 1:3 are ALR coords of the deposit/secondary/post energy
            # simplex; decode them against pre_E so edep + e_sec + post_E == pre_E
            # (hence delta_e == edep + e_sec) holds by construction. e_sec_pred
            # doubles as the stick-breaking energy budget for the Stage-2 decode
            # below, since the model has no other source for it at inference.
            edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(
                raw[:, 1:3], piece["pre_E"]
            )

            # Normalise predicted direction then rotate back to world frame
            post_dir_local = raw[:, 3:6].copy()
            norms = np.linalg.norm(post_dir_local, axis=1, keepdims=True)
            post_dir_local /= np.where(norms < 1e-8, 1.0, norms)
            post_dir_world = inv_local_frame_rotation(piece["pre_dir"], post_dir_local)

            # Same for the travel direction, then reconstruct post_pos from
            # the single shared step_length so the two stay consistent.
            travel_dir_local = raw[:, 6:9].copy()
            norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True)
            travel_dir_local /= np.where(norms < 1e-8, 1.0, norms)
            post_pos_world = reconstruct_post_pos(
                piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
            )

            sec_E, sec_dir_world, sec_pdg_code, _sec_valid = decode_secondaries(
                sec_cont_np,
                sec_pdg_idx_np,
                n_sec_pred_np,
                e_sec_pred,
                piece["pre_dir"],
                pdg_map_inv,
            )
            sec_pdg_list = [
                sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
            ]
            sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
            sec_dx_list = [
                sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)
            ]
            sec_dy_list = [
                sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)
            ]
            sec_dz_list = [
                sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)
            ]

            table = pa.table(
                {
                    "event_id": piece["event_id"],
                    "pdg": piece["pdg"],
                    "pre_x": piece["pre_pos"][:, 0],
                    "pre_y": piece["pre_pos"][:, 1],
                    "pre_z": piece["pre_pos"][:, 2],
                    "pre_E": piece["pre_E"],
                    "pre_dx": piece["pre_dir"][:, 0],
                    "pre_dy": piece["pre_dir"][:, 1],
                    "pre_dz": piece["pre_dir"][:, 2],
                    "material": piece["material"],
                    "layer_id": piece["layer_id"],
                    "n_sec": piece["n_sec"],
                    "n_sec_pred": n_sec_pred_np,
                    "step_length": step_length,
                    "delta_e": delta_e,
                    "edep": edep,
                    "post_dx": post_dir_world[:, 0],
                    "post_dy": post_dir_world[:, 1],
                    "post_dz": post_dir_world[:, 2],
                    "post_x": post_pos_world[:, 0],
                    "post_y": post_pos_world[:, 1],
                    "post_z": post_pos_world[:, 2],
                    "sec_pdg_list": sec_pdg_list,
                    "sec_E_list": sec_E_list,
                    "sec_dx_list": sec_dx_list,
                    "sec_dy_list": sec_dy_list,
                    "sec_dz_list": sec_dz_list,
                }
            )

        table = table.replace_schema_metadata(
            {
                PREDICT_COORD_METADATA_KEY: coord.value,
                PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
            }
        )

        if writer is None:
            writer = pq.ParquetWriter(out, table.schema)
        writer.write_table(table)
        total += len(piece["event_id"])

    # Buffer rows across row-group boundaries so the inference batch size
    # isn't capped by however the source file happens to be chunked.
    buffer: dict[str, np.ndarray] | None = None

    bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
    for path in files:
        for chunk in chunk_iter(path):
            N_in = len(chunk["event_id"])

            pdg_mask = np.array([int(p) in pdg_map for p in chunk["pdg"]])
            if not pdg_mask.all():
                unknown_pdg_counts.update(int(p) for p in chunk["pdg"][~pdg_mask])
                chunk = {k: v[pdg_mask] for k, v in chunk.items()}

            skipped += N_in - len(chunk["event_id"])
            bar.update(N_in)
            if len(chunk["event_id"]) == 0:
                continue

            buffer = chunk if buffer is None else _concat(buffer, chunk)
            while len(buffer["event_id"]) >= bs:
                piece = {k: v[:bs] for k, v in buffer.items()}
                buffer = {k: v[bs:] for k, v in buffer.items()}
                _process(piece)

    if buffer is not None and len(buffer["event_id"]) > 0:
        _process(buffer)

    bar.close()
    if writer is not None:
        writer.close()

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment)
    typer.echo(f"reference: {ref_path}")

    if skipped:
        codes = ", ".join(
            f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items())
        )
        typer.echo(
            f"warning: skipped {skipped:,} row(s) with unknown PDG code(s): {codes}",
            err=True,
        )
    typer.echo(f"wrote {total:,} rows → {out}")


def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.ndarray]:
    """Pick each event's primary entry state (argmax-pre_E row) as a shower seed.

    Streams conditioning columns and keeps the highest-pre_E step per event_id —
    the codebase's convention for the primary (a secondary always carries less
    energy than its parent). See giant/analysis.py:_entry_axis_and_bin_edges.
    """
    best_E: dict[int, float] = {}
    best: dict[int, tuple] = {}
    for path in files:
        for chunk in iter_cond_chunks(path):
            ev = chunk["event_id"]
            pe = chunk["pre_E"]
            for i in range(len(ev)):
                e = int(ev[i])
                if pe[i] > best_E.get(e, -np.inf):
                    best_E[e] = float(pe[i])
                    best[e] = (
                        int(chunk["pdg"][i]),
                        chunk["pre_pos"][i].astype(np.float64),
                        float(pe[i]),
                        chunk["pre_dir"][i].astype(np.float64),
                    )
    event_ids = sorted(best)
    if n_events is not None:
        event_ids = event_ids[:n_events]
    if not event_ids:
        raise ValueError("no events found to seed from")

    return {
        "event_id": np.array(event_ids, dtype=np.int64),
        "pdg": np.array([best[e][0] for e in event_ids], dtype=np.int64),
        "pre_pos": np.stack([best[e][1] for e in event_ids]),
        "pre_E": np.array([best[e][2] for e in event_ids], dtype=np.float64),
        "pre_dir": np.stack([best[e][3] for e in event_ids]),
    }


@app.command()
def rollout(
    data: Annotated[
        Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")
    ],
    checkpoint: Annotated[
        Path,
        typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"),
    ],
    geometry: Annotated[
        Path,
        typer.Option(
            "--geometry",
            "-g",
            help="Geometry oracle .pkl (dwarf build-geometry-oracle)",
        ),
    ],
    energy_cutoff: Annotated[
        float,
        typer.Option(
            "--energy-cutoff",
            help="Stop a track when its energy drops below this [MeV]",
        ),
    ] = 0.1,
    max_steps: Annotated[
        int, typer.Option("--max-steps", help="Max steps per individual track")
    ] = 1000,
    steps: Annotated[
        int,
        typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"),
    ] = 10,
    batch_size: Annotated[
        int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")
    ] = 4096,
    max_tracks_per_event: Annotated[
        Optional[int],
        typer.Option(
            "--max-tracks-per-event",
            help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)",
        ),
    ] = None,
    escape_threshold: Annotated[
        Optional[float],
        typer.Option(
            "--escape-threshold",
            help="Override the oracle's NN-distance escape threshold [mm]",
        ),
    ] = None,
    n_events: Annotated[
        Optional[int], typer.Option("--n-events", help="Cap number of seed events")
    ] = None,
    device: Annotated[
        Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")
    ] = None,
    out: Annotated[
        Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")
    ] = None,
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
    ] = None,
) -> None:
    """Roll the surrogate forward into full showers (autoregressive)."""
    if seed is not None:
        torch.manual_seed(seed)
        np.random.seed(seed)

    _device = torch.device(device) if device else gconfig.auto_device()
    typer.echo(f"device: {_device}")

    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
    for key in ("model_config", "sec_decoder"):
        if key not in ckpt:
            typer.echo(
                f"error: checkpoint has no {key} — retrain with the current code",
                err=True,
            )
            raise typer.Exit(1)

    model_cfg = ckpt["model_config"]
    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])

    model, sec_decoder = build_models(model_cfg)
    model.load_state_dict(ckpt["model"])
    model.to(_device).eval()
    sec_decoder.load_state_dict(ckpt["sec_decoder"])
    sec_decoder.to(_device).eval()
    typer.echo(f"loaded checkpoint: {checkpoint}")

    oracle = GeometryOracle.load(geometry)
    typer.echo(
        f"loaded geometry oracle: {geometry} "
        f"(escape_threshold={oracle.escape_threshold:.3f})"
    )

    files = find_parquet_files(data)
    seeds = _seed_from_data(files, n_events)
    typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")

    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
    out.parent.mkdir(parents=True, exist_ok=True)

    # Written incrementally as each batch of steps is produced, rather than
    # buffering the whole run (which scales with n_events * max_steps *
    # avg_tracks_per_event) — mirrors the row-group streaming `giant predict`
    # already does on its input side.
    writer: pq.ParquetWriter | None = None

    def _write_chunk(row: dict[str, np.ndarray]) -> None:
        nonlocal writer
        table = pa.table(row)
        if writer is None:
            table = table.replace_schema_metadata(
                {
                    PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE,
                    PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
                }
            )
            writer = pq.ParquetWriter(out, table.schema)
        writer.write_table(table)

    summary = run_rollout(
        model,
        sec_decoder,
        oracle,
        seeds,
        cond_norm,
        tgt_norm,
        pdg_map,
        mat_map,
        energy_cutoff=energy_cutoff,
        max_steps=max_steps,
        steps=steps,
        batch_size=batch_size,
        device=_device,
        max_tracks_per_event=max_tracks_per_event,
        escape_threshold=escape_threshold,
        on_chunk=_write_chunk,
    )
    if writer is not None:
        writer.close()

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
    ref = yaml.safe_load(ref_path.read_text())
    ref.update(
        {
            "kind": "rollout",
            "geometry_oracle": str(geometry.resolve()),
            "energy_cutoff": energy_cutoff,
            "max_steps": max_steps,
            "steps": steps,
            "max_tracks_per_event": max_tracks_per_event,
            "n_seed_events": int(len(seeds["event_id"])),
        }
    )
    ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))

    typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
    typer.echo(f"terminations: {summary['termination_reason_counts']}")
    typer.echo(f"reference: {ref_path}")


if __name__ == "__main__":
    app()