giant/giant/data/loader.py

from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

import numpy as np
import pandas as pd
import pyarrow.parquet as pq

from giant.constants import K_MAX

# A manifest is a plain text file listing one parquet path per line, used to
# name a curated subset of files (e.g. a train/holdout pool) without copying
# or symlinking the underlying parquet files. Lines are resolved relative to
# the manifest's own directory, so the manifest stays valid if the whole
# dataset tree is moved or copied elsewhere intact.
MANIFEST_SUFFIX = ".manifest"

# Each input parquet file is a separate Geant4 job converted 1:1 from its own
# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
# always restarts from 0 — so when multiple files are loaded together (a
# directory or .manifest), raw event_id values collide across files even
# though they refer to unrelated events. Every per-file event_id column gets
# offset by its file's index in the (deterministically ordered) files list
# so ids stay globally unique across a multi-file load; the stride is far
# larger than any realistic per-file event count.
EVENT_ID_FILE_STRIDE = 1_000_000


def event_id_offset(file_index: int) -> int:
    return file_index * EVENT_ID_FILE_STRIDE


def _offset_event_id(raw_ids: np.ndarray, offset: int) -> np.ndarray:
    """Add this file's `event_id_offset`, after checking the raw ids fit in one stride block.

    Without this check, a file whose own raw event_id numbering reaches
    `EVENT_ID_FILE_STRIDE` (an unusually large job, or non-contiguous
    numbering) would silently collide into the next file's offset block,
    merging unrelated events across files — reintroducing exactly the
    train/val event leakage this offset scheme exists to prevent.
    """
    raw_ids = np.asarray(raw_ids, dtype=np.int64)
    if raw_ids.size and int(raw_ids.max()) >= EVENT_ID_FILE_STRIDE:
        raise ValueError(
            f"event_id {int(raw_ids.max())} >= EVENT_ID_FILE_STRIDE "
            f"({EVENT_ID_FILE_STRIDE}) — this file has a larger event_id "
            "than the per-file offset scheme can support without colliding "
            "with the next file's id block."
        )
    return raw_ids + offset


def _read_manifest(path: Path) -> list[Path]:
    files = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        resolved = (path.parent / line).resolve()
        if not resolved.is_file():
            raise FileNotFoundError(f"{path} lists missing file: {resolved}")
        files.append(resolved)
    if not files:
        raise FileNotFoundError(f"manifest {path} lists no files")
    return files


def find_parquet_files(path: str | Path) -> list[Path]:
    p = Path(path)
    if p.suffix == MANIFEST_SUFFIX:
        return _read_manifest(p)
    if p.is_dir():
        files = sorted(p.glob("*.parquet"))
        if not files:
            raise FileNotFoundError(f"no .parquet files found in {p}")
        return files
    return [p]


def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
    """Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
    out = np.full((len(series), K), fill, dtype=np.float32)
    for i, lst in enumerate(series):
        if lst is not None and len(lst) > 0:
            n = min(len(lst), K)
            out[i, :n] = lst[:n]
    return out


def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
    """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
    out = np.full((len(series), K), fill, dtype=np.int64)
    for i, lst in enumerate(series):
        if lst is not None and len(lst) > 0:
            n = min(len(lst), K)
            out[i, :n] = lst[:n]
    return out


def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
    """Pad three list-valued direction columns → (N, K, 3) float32.

    Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
    """
    N = len(dx)
    out = np.zeros((N, K, 3), dtype=np.float32)
    out[:, :, 2] = 1.0
    for i in range(N):
        lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
        if lx is not None and len(lx) > 0:
            n = min(len(lx), K)
            out[i, :n, 0] = lx[:n]
            out[i, :n, 1] = ly[:n]
            out[i, :n, 2] = lz[:n]
    return out


def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
    has_sec_lists = "sec_E_list" in df.columns

    d: dict[str, np.ndarray] = {
        "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
        "pdg": df["pdg"].to_numpy(dtype=np.int32),
        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
        "material": df["material"].to_numpy(dtype=object),
        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
        # The physics process that ended the step (e.g. "compt", "phot",
        # "eBrem") — a post-step outcome, so it's a router/classifier
        # supervision label only, never conditioning (see build_process_map*
        # / ProcessRouter). Guarded like has_sec_lists: older parquet
        # conversions predating this column still load fine.
        "process": (
            df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
        ),
        "step_length": df["step_length"].to_numpy(dtype=np.float32),
        "post_E": df["post_E"].to_numpy(dtype=np.float32),
        "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
        "edep": df["edep"].to_numpy(dtype=np.float32),
        "post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
        "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
    }

    if has_sec_lists:
        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
        d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)

    return d


def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
    return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)


def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
    """Read only the event_id column — cheap scan for split assignment."""
    ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
    return _offset_event_id(ids, offset)


def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> Iterator[dict[str, np.ndarray]]:
    """Yield one parquet row-group at a time so a large file never fully loads.

    `k_max` sets the padded width of the sec_*_list columns (should match
    `stage2_model.k_max`); defaults to the
    module constant for callers that don't care (e.g. Stage-1-only reads)."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max)


_COND_COLS = [
    "event_id",
    "pdg",
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "material",
    "layer_id",
    "child_track_ids",
    "e_sec",
]


def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
    return {
        "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
        "pdg": df["pdg"].to_numpy(dtype=np.int32),
        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
        "material": df["material"].to_numpy(dtype=object),
        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
    }


def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np.ndarray]]:
    """Yield conditioning-only row-groups (no post-step columns read from disk)."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)


def build_index_maps(
    data: dict[str, np.ndarray],
) -> tuple[dict[int, int], dict[str, int]]:
    pdg_vals = sorted(int(v) for v in np.unique(data["pdg"]))
    mat_vals = sorted(str(v) for v in np.unique(data["material"]))
    return (
        {v: i for i, v in enumerate(pdg_vals)},
        {v: i for i, v in enumerate(mat_vals)},
    )


def build_index_maps_from_files(
    files: list[Path],
) -> tuple[dict[int, int], dict[str, int]]:
    """Scan only pdg and material columns across all files (2-column read)."""
    pdg_vals: set[int] = set()
    mat_vals: set[str] = set()
    for path in files:
        df = pd.read_parquet(path, columns=["pdg", "material"])
        pdg_vals.update(int(v) for v in df["pdg"].unique())
        mat_vals.update(str(v) for v in df["material"].unique())
    return (
        {v: i for i, v in enumerate(sorted(pdg_vals))},
        {v: i for i, v in enumerate(sorted(mat_vals))},
    )


def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None:
    for name, count in series.value_counts().items():
        name = cast(name)
        counts[name] = counts.get(name, 0) + int(count)


def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
    """Scan `column` across `files` and return `{cast(value): total_count}`,
    accumulated in file order (see `fingerprint_files`'s docstring on why
    scan order — not a normalized/sorted order — is preserved: it drives
    tie-breaking in the frequency ranking below)."""
    counts: dict = {}
    for path in files:
        df = pd.read_parquet(path, columns=[column])
        _accumulate_value_counts(counts, df[column], cast)
    return counts


def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]:
    """Frequency-capped value->index map: the `n_classes - 1` most frequent
    keys get their own index; every rarer key is bucketed into a shared
    "other" index (`n_classes - 1`).

    Returns `(class_map, other_members)` — `other_members` is `{key: count}`
    for every key bucketed into "other" (the empirical within-bucket
    distribution, for `other_policy = "sample"` at rollout).
    """
    ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
    keep = ranked[: max(n_classes - 1, 0)]
    class_map = {k: i for i, k in enumerate(keep)}
    other_idx = n_classes - 1
    other_members: dict = {}
    for k in ranked[len(keep) :]:
        class_map[k] = other_idx
        other_members[k] = counts[k]
    return class_map, other_members


def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
    """Scan the `process` column and build a frequency-capped process->index map.

    Physics processes have a long tail (rare nuclear captures, decays, ...)
    while `ProcessRouter` needs a fixed number of expert slots, so only the
    `n_experts - 1` most frequent processes get their own index; every rarer
    process is bucketed into a shared "other" index (`n_experts - 1`). This
    mirrors how `build_features` clamps the n_sec label to K_MAX for the
    fixed-width n_sec_head classifier.
    """
    counts = _rank_by_frequency_from_files(files, "process", str)
    class_map, _ = _topn_plus_other_map(counts, n_experts)
    return class_map


@dataclass
class TopNMap:
    """A frequency-capped value->index map for a conditioning/type axis (PDG
    or material), plus the empirical within-bucket distribution of whatever
    got folded into "other" — see `build_topn_map_from_files`."""

    class_map: dict
    other_members: dict


def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
    """Scan `column` and build a frequency-capped value->index map, structurally
    identical to `build_process_map_from_files` (shares its ranking core via
    `_topn_plus_other_map`), generalized over the source column and key type.

    Used for the material axis (`column="material"`, `cast=str`, matching
    `mat_map`'s key type). The PDG axis uses
    `build_pdg_topn_map_from_files` instead (it needs to pool two columns,
    which this single-column form can't express). Also records
    `other_members` (the empirical within-"other" distribution), needed
    later for `other_policy = "sample"` at rollout — computed now since it's
    free during this same scan.
    """
    counts = _rank_by_frequency_from_files(files, column, cast)
    class_map, other_members = _topn_plus_other_map(counts, n_classes)
    return TopNMap(class_map=class_map, other_members=other_members)


def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
    """PDG top-N-plus-other map, pooling counts from BOTH roles a PDG code
    plays in this dataset: a step's own primary particle (`pdg` column) and
    an emitted secondary's species (`sec_pdg_list`, exploded) — shared by
    `conditioning.particle.type = "onehot"` and
    `stage2_model.particle_type.target = "onehot"`. Pooling both is what
    keeps a species that's common as a secondary
    but rare as a primary (or vice versa) from being pushed into "other"
    just because one role's count alone looks small — the meeting's failure
    mode (zero photon secondaries, hallucinated antineutrinos) was
    specifically about secondary-species collapse, so the map this feeds
    needs to reflect secondary frequency, not just primary frequency.

    `sec_pdg_list` is absent from parquet files predating the parent->child
    join (see `_df_to_dict`'s `has_sec_lists` guard) — silently skipped for
    those, same convention as elsewhere in this module.
    """
    counts: dict = {}
    for path in files:
        columns = ["pdg"]
        has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names
        if has_sec:
            columns.append("sec_pdg_list")
        df = pd.read_parquet(path, columns=columns)
        _accumulate_value_counts(counts, df["pdg"], int)
        if has_sec:
            exploded = df["sec_pdg_list"].explode().dropna()
            _accumulate_value_counts(counts, exploded, int)
    class_map, other_members = _topn_plus_other_map(counts, n_classes)
    return TopNMap(class_map=class_map, other_members=other_members)