giant/giant/data/transforms.py

import warnings

import numpy as np

_EPS = 1e-8

# Floor added to each energy fraction before taking log-ratios so the simplex
# coordinates stay finite on the heavily populated boundary: e_sec is 0 in ~half
# of all steps and post_E is 0 at every track end. Kept tiny (1e-5 of pre_E) so
# the conservation it slightly softens is physically negligible (~0.001%).
_SIMPLEX_FLOOR = 1e-5


def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
    return np.log(np.asarray(x, dtype=np.float32) + eps)


def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
    return np.exp(np.asarray(y, dtype=np.float32)) - eps


def energy_simplex_encode(
    edep: np.ndarray,
    e_sec: np.ndarray,
    post_E: np.ndarray,
    pre_E: np.ndarray,
) -> np.ndarray:
    """Encode (edep, e_sec, post_E) as 2 additive-log-ratio (ALR) coordinates.

    The three energies are first expressed as fractions of pre_E that sum to 1:
    post_E is kept exact (so the particle's retained energy — the next step's
    input during a rollout — is preserved), and the energy actually lost,
    `delta_e = pre_E - post_E`, is split between local deposit and secondaries
    in the recorded edep:e_sec ratio (re-attributing any sub-threshold / rest-mass
    leakage proportionally so the three fractions sum to exactly 1). A small floor
    (`_SIMPLEX_FLOOR`) keeps the log-ratios finite where a fraction is 0.

    Returns an (N, 2) array of ALR coordinates referenced to the post fraction;
    `energy_simplex_decode` is its inverse (up to the floor softening).
    """
    pre_E = np.maximum(np.asarray(pre_E, dtype=np.float32), _EPS)
    raw_post_E = np.asarray(post_E, dtype=np.float32)
    post_E = np.clip(raw_post_E, 0.0, pre_E)
    delta_e = pre_E - post_E
    edep = np.asarray(edep, dtype=np.float32)
    e_sec = np.asarray(e_sec, dtype=np.float32)
    lost = edep + e_sec
    has_loss = lost > _EPS
    # Clipping post_E down to pre_E forces delta_e (and thus the rescaled
    # edep/e_sec below) to 0 even on rows where edep/e_sec were genuinely
    # recorded as nonzero — warn so this doesn't silently discard real data.
    discarded = (raw_post_E > pre_E) & has_loss
    if np.any(discarded):
        n = int(np.sum(discarded))
        warnings.warn(
            f"energy_simplex_encode: {n}/{len(discarded)} step(s) had "
            "post_E > pre_E (clipped) while recording nonzero edep/e_sec; "
            "that recorded energy deposit is discarded to keep delta_e "
            "consistent with the clip.",
            stacklevel=2,
        )
    scale = np.where(has_loss, delta_e / np.maximum(lost, _EPS), 0.0)
    # Where nothing was recorded as deposited/secondary but energy was lost,
    # attribute all of delta_e to local deposit.
    edep_r = np.where(has_loss, edep * scale, delta_e)
    e_sec_r = np.where(has_loss, e_sec * scale, 0.0)

    f = np.stack([edep_r, e_sec_r, post_E], axis=-1) / pre_E[:, None]  # (N,3), sums≈1
    f = (f + _SIMPLEX_FLOOR) / (1.0 + 3.0 * _SIMPLEX_FLOOR)
    log_f = np.log(f)
    z = log_f[:, :2] - log_f[:, 2:3]  # ALR vs post reference
    return z.astype(np.float32)


def energy_simplex_decode(
    z: np.ndarray, pre_E: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.

    A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
    `edep + e_sec + post_E == pre_E` holds by construction (the energy-conservation
    inductive bias). Returns `(edep, e_sec, post_E, delta_e)` in physical units.
    """
    z = np.asarray(z, dtype=np.float32)
    pre_E = np.asarray(pre_E, dtype=np.float32)
    logits = np.concatenate([z, np.zeros((len(z), 1), dtype=np.float32)], axis=1)
    logits = logits - logits.max(axis=1, keepdims=True)
    f = np.exp(logits)
    f /= f.sum(axis=1, keepdims=True)  # sums to 1 exactly → exact conservation
    E = f * pre_E[:, None]
    edep, e_sec, post_E = E[:, 0], E[:, 1], E[:, 2]
    delta_e = pre_E - post_E
    return (
        edep.astype(np.float32),
        e_sec.astype(np.float32),
        post_E.astype(np.float32),
        delta_e.astype(np.float32),
    )


def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
    """Unit rotation axis `pre_dir × ẑ`, closed-form since ẑ = [0, 0, 1] is constant.

    `cross(a, [0,0,1]) = [a_y, -a_x, 0]` — substituting the constant operand
    avoids a generic `np.cross` call (shape/broadcast handling for an
    arbitrary second operand) on every row; profiling on a 114M-row file
    showed `np.cross` as the single hottest call inside this rotation.
    """
    axis = np.stack(
        [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
    )
    axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)  # (N,1)
    # Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
    # are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
    safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
    return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)


def _cross_with_z_axis(axis: np.ndarray, v: np.ndarray) -> np.ndarray:
    """`axis × v`, closed-form since `axis` from `_rodrigues_axis` always has z = 0.

    `cross([ax,ay,0], [bx,by,bz]) = [ay*bz, -ax*bz, ax*by - ay*bx]`.
    """
    ax, ay = axis[:, 0:1], axis[:, 1:2]
    bx, by, bz = v[:, 0:1], v[:, 1:2], v[:, 2:3]
    return np.concatenate([ay * bz, -ax * bz, ax * by - ay * bx], axis=1)


def _validate_unit_pre_dir(pre_dir: np.ndarray) -> np.ndarray:
    """Normalize pre_dir and raise if any row is too degenerate to define a frame.

    `local_frame_rotation`/`inv_local_frame_rotation` treat pre_dir[:, 2] as
    cos(angle to ẑ), which is only correct for a unit vector. Small float32
    drift is corrected silently; a near-zero-norm row has no well-defined
    direction, so it's raised loudly instead of producing a meaningless
    rotation (previously it fell through to an arbitrary axis with no error).
    """
    pre_dir = np.asarray(pre_dir, dtype=np.float32)
    norm = np.linalg.norm(pre_dir, axis=1, keepdims=True)
    if np.any(norm < 1e-6):
        raise ValueError(
            f"pre_dir has {int(np.sum(norm < 1e-6))} row(s) with near-zero norm "
            "(< 1e-6); local/inv_local_frame_rotation require a well-defined "
            "incoming direction for every row."
        )
    return pre_dir / norm


def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
    """Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).

    Preserves the angle between pre_dir and post_dir; the result has post_dir
    expressed relative to a coordinate system in which the incoming particle
    travels along +z.
    """
    pre_dir = _validate_unit_pre_dir(pre_dir)
    cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0)  # (N,1); dot with ẑ = z-component
    sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2))  # (N,1)

    axis = _rodrigues_axis(pre_dir)  # (N,3); zero-z, zero-norm when pre_dir ∥ ẑ
    kxv = _cross_with_z_axis(axis, post_dir)  # (N,3)
    kdv = (axis * post_dir).sum(axis=1, keepdims=True)  # (N,1)

    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
        np.float32
    )


class Normalizer:
    def __init__(self) -> None:
        self.mean: np.ndarray | None = None
        self.std: np.ndarray | None = None

    def fit(self, X: np.ndarray) -> "Normalizer":
        self.mean = X.mean(axis=0).astype(np.float32)
        self.std = X.std(axis=0).astype(np.float32)
        self.std = np.where(self.std < _EPS, 1.0, self.std).astype(np.float32)
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        return ((X - self.mean) / self.std).astype(np.float32)

    def inverse_transform(self, X: np.ndarray) -> np.ndarray:
        return (X * self.std + self.mean).astype(np.float32)

    def to_dict(self) -> dict:
        assert self.mean is not None and self.std is not None, "Normalizer not fitted"
        return {"mean": self.mean.tolist(), "std": self.std.tolist()}

    @classmethod
    def from_dict(cls, d: dict) -> "Normalizer":
        obj = cls()
        obj.mean = np.array(d["mean"], dtype=np.float32)
        obj.std = np.array(d["std"], dtype=np.float32)
        return obj


class _WelfordAccumulator:
    """Streaming mean/variance (Welford's online algorithm, batch update).

    Use to fit a Normalizer over data that doesn't fit in memory:
        acc = _WelfordAccumulator(n_features)
        for chunk in data:
            acc.update(chunk)
        normalizer = acc.to_normalizer()
    """

    def __init__(self, n_features: int) -> None:
        self.n = 0
        self._mean = np.zeros(n_features, dtype=np.float64)
        self._M2 = np.zeros(n_features, dtype=np.float64)

    def update(self, X: np.ndarray) -> None:
        X = np.asarray(X, dtype=np.float64)
        B = X.shape[0]
        new_n = self.n + B
        delta = X - self._mean
        self._mean += delta.sum(0) / new_n
        delta2 = X - self._mean
        self._M2 += (delta * delta2).sum(0)
        self.n = new_n

    def to_normalizer(self) -> "Normalizer":
        norm = Normalizer()
        norm.mean = self._mean.astype(np.float32)
        std = np.sqrt(self._M2 / max(self.n, 1)).astype(np.float32)
        norm.std = np.where(std < _EPS, 1.0, std).astype(np.float32)
        return norm


class _ReservoirSampler:
    """Uniform random sample of a fixed capacity drawn from a data stream.

    Algorithm R (Vitter 1985), vectorized per chunk so it stays cheap over
    hundreds of millions of rows: use to get a representative subsample of
    a column for a distribution estimate (e.g. quantiles) without
    materializing the full column.

        sampler = _ReservoirSampler(capacity=100_000)
        for chunk in data:
            sampler.update(chunk)
        sample = sampler.sample
    """

    def __init__(self, capacity: int, seed: int = 0) -> None:
        self.capacity = capacity
        self.n_seen = 0
        self._rng = np.random.default_rng(seed)
        self._reservoir = np.empty(0, dtype=np.float64)

    def update(self, values: np.ndarray) -> None:
        values = np.asarray(values, dtype=np.float64).reshape(-1)
        if values.size == 0:
            return
        n_before = self.n_seen
        if n_before < self.capacity:
            take = min(values.size, self.capacity - n_before)
            self._reservoir = np.concatenate([self._reservoir, values[:take]])
            values = values[take:]
            n_before += take
        self.n_seen = n_before + values.size
        if values.size == 0 or self.capacity == 0:
            return
        # remaining elements are past the fill phase: element at 1-based
        # stream position j replaces a uniformly random reservoir slot with
        # probability capacity/j, which yields a uniform sample overall.
        positions = n_before + np.arange(1, values.size + 1)
        accept = self._rng.random(values.size) < (self.capacity / positions)
        accept_idx = np.nonzero(accept)[0]
        if accept_idx.size > 0:
            slots = self._rng.integers(0, self.capacity, size=accept_idx.size)
            self._reservoir[slots] = values[accept_idx]

    @property
    def sample(self) -> np.ndarray:
        return self._reservoir.astype(np.float32)


def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
    """World-frame unit vector pointing from pre_pos to post_pos.

    Kept independent of `step_length`: that scalar already encodes the
    magnitude of this displacement, so this function only ever returns
    direction (norm-guarded the same way as `local_frame_rotation`'s axis).
    """
    disp = post_pos - pre_pos
    norm = np.linalg.norm(disp, axis=1, keepdims=True)
    safe_norm = np.where(norm < 1e-7, 1.0, norm)
    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
        np.float32
    )


def reconstruct_post_pos(
    pre_pos: np.ndarray,
    pre_dir: np.ndarray,
    step_length: np.ndarray,
    travel_dir_local: np.ndarray,
) -> np.ndarray:
    """Inverse of the travel_direction/local_frame_rotation encoding.

    Single source of truth for combining the magnitude (`step_length`) and
    direction (`travel_dir_local`) back into a world-frame post_pos, so
    `step_length` and post_pos stay consistent by construction.
    """
    travel_dir_world = inv_local_frame_rotation(pre_dir, travel_dir_local)
    return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)


def inv_local_frame_rotation(
    pre_dir: np.ndarray, post_dir_local: np.ndarray
) -> np.ndarray:
    """Inverse of local_frame_rotation: rotate from local frame back to world frame.

    Applies R^T (same axis, negative angle) to post_dir_local.
    """
    pre_dir = _validate_unit_pre_dir(pre_dir)
    cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0)  # dot with ẑ = z-component
    sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2))

    axis = _rodrigues_axis(pre_dir)
    kxv = _cross_with_z_axis(axis, post_dir_local)
    kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)

    # Negative angle: sin_t → -sin_t
    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
        np.float32
    )


_STICK_LOGIT_CLIP = 10.0  # logit value used for the last valid secondary slot


def encode_secondaries(
    sec_E_list: np.ndarray,
    sec_dir_list: np.ndarray,
    sec_valid: np.ndarray,
    e_sec: np.ndarray,
    pre_dir: np.ndarray,
    sec_pdg_list: np.ndarray | None = None,
) -> np.ndarray:
    """Encode per-secondary attributes into continuous per-slot targets.

    Secondaries must already be sorted descending by energy (as stored in the
    parquet). Returns sec_cont of shape (N, K_MAX, SEC_SLOT_DIM=6):
      slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z,
                 log_mass, charge]

    Stick-breaking logit: for slot i, f_i = E_i / remaining_budget, where
    remaining_budget = e_sec - sum(E_0..E_{i-1}). The logit is log(f/(1-f)),
    clipped to ±_STICK_LOGIT_CLIP. The last valid slot gets +_STICK_LOGIT_CLIP
    (takes the full remaining budget). Padding slots get 0.

    log_mass/charge are the secondary's real physical identity, looked up
    from its ground-truth PDG code (`sec_pdg_list`) via
    `giant.particles.particle_phys_array` — a fixed physics-derived
    regression target, not a learned/moving one (unlike the embedding-table
    target this replaced), so nothing needs to be detached at training time.
    `sec_pdg_list` is optional so callers that only need the continuous
    stick/dir block (e.g. inference-time re-encoding) can omit it; omitting
    it zero-fills the last two columns, matching the padding-slot convention.
    """
    N, K = sec_E_list.shape
    e_sec = np.asarray(e_sec, dtype=np.float64)

    stick_logits = np.zeros((N, K), dtype=np.float32)
    for i in range(K):
        if i == 0:
            remaining = np.maximum(e_sec, _EPS)
        else:
            remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS)
        f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
        logit = np.log(f / (1.0 - f)).astype(np.float32)
        # Last valid slot: give it the full remaining budget
        is_last = sec_valid[:, i] & ~(
            sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
        )
        logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
        logit = np.where(
            sec_valid[:, i], np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), 0.0
        )
        stick_logits[:, i] = logit.astype(np.float32)

    # Rotate each slot's direction into the local frame of the primary.
    # pre_dir is broadcast across all K slots.
    dir_local = np.zeros((N, K, 3), dtype=np.float32)
    for i in range(K):
        # Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
        valid_mask = sec_valid[:, i]
        if valid_mask.any():
            dir_local[valid_mask, i] = local_frame_rotation(
                pre_dir[valid_mask], sec_dir_list[valid_mask, i]
            )

    if sec_pdg_list is not None:
        from giant.particles import particle_phys_array

        # Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int),
        # which isn't a resolvable particle — substitute a dummy resolvable
        # code (22, photon) there, since the result is discarded below by
        # the sec_valid mask regardless.
        safe_pdg = np.where(sec_valid, sec_pdg_list, 22)
        flat_mass_charge = particle_phys_array(safe_pdg.reshape(-1))  # (N*K, 2)
        mass = flat_mass_charge[:, 0].reshape(N, K)
        charge = flat_mass_charge[:, 1].reshape(N, K)
        log_mass = np.where(sec_valid, log_transform(mass), 0.0).astype(np.float32)
        charge = np.where(sec_valid, charge, 0.0).astype(np.float32)
    else:
        log_mass = np.zeros((N, K), dtype=np.float32)
        charge = np.zeros((N, K), dtype=np.float32)

    sec_cont = np.concatenate(
        [stick_logits[:, :, None], dir_local, log_mass[:, :, None], charge[:, :, None]],
        axis=-1,
    )  # (N, K, 6)
    return sec_cont.astype(np.float32)


def decode_secondaries(
    sec_cont: np.ndarray,
    n_sec: np.ndarray,
    e_sec: np.ndarray,
    pre_dir: np.ndarray,
    sec_phys_normalizer: "Normalizer | None" = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Inverse of encode_secondaries: continuous targets → physical secondary attrs.

    sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y,
        local_dir_z, log_mass, charge] (log_mass/charge normalised iff
        `sec_phys_normalizer` was applied when this was produced — e.g. a
        raw model prediction; pass the same normalizer here to invert it)
    n_sec: (N,) integer secondary counts
    e_sec: (N,) total secondary energy budget [MeV]
    pre_dir: (N, 3) pre-step world-frame direction

    Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each
    shape (N, K_MAX). The valid slots' energies (`sec_E[sec_valid]`, per row)
    always sum to exactly `e_sec` — see the rescaling below. mass/charge are
    the model's raw predicted physical identity for each secondary, used
    as-is (no snapping to a discrete PDG code) — see giant/particles.py for
    the separate, reporting-only nearest-PDG lookup callers may apply on top
    of this for display/bookkeeping purposes.
    """
    if sec_phys_normalizer is not None:
        N_, K_, _ = sec_cont.shape
        phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2))
        sec_cont = sec_cont.copy()
        sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)

    N, K, _ = sec_cont.shape
    stick_logits = sec_cont[:, :, 0]  # (N, K)
    dir_local = sec_cont[:, :, 1:4].copy()  # (N, K, 3)
    log_mass = sec_cont[:, :, 4]  # (N, K)
    charge = sec_cont[:, :, 5]  # (N, K)

    # Flow-matching output isn't guaranteed unit norm; normalise before the
    # rotation below, which preserves magnitude rather than fixing it up.
    norms = np.linalg.norm(dir_local, axis=-1, keepdims=True)
    dir_local /= np.where(norms < 1e-8, 1.0, norms)

    fractions = 1.0 / (1.0 + np.exp(-stick_logits.astype(np.float64)))

    sec_E = np.zeros((N, K), dtype=np.float64)
    e_sec = np.asarray(e_sec, dtype=np.float64)
    remaining = e_sec.copy()
    for i in range(K):
        sec_E[:, i] = fractions[:, i] * remaining
        remaining = np.maximum(remaining - sec_E[:, i], 0.0)

    sec_valid = np.arange(K)[None, :] < n_sec[:, None]  # (N, K)

    # Stick-breaking guarantees sum(sec_E[valid]) <= e_sec (each fraction is in
    # [0,1] of an already-shrinking remainder) but rarely hits it exactly, so
    # rescale the valid slots by one common per-row factor to close that gap —
    # rather than dumping the shortfall into whichever slot happens to be last
    # by energy rank, which would let one low-energy secondary balloon and
    # distort the shower's topology. This preserves each row's relative split
    # across its secondaries and only ever scales up (valid_sum <= e_sec).
    # Rows where every valid slot decoded to ~zero (scale undefined) fall back
    # to an even split of e_sec across the n_sec valid slots.
    sec_E = sec_E * sec_valid
    valid_sum = sec_E.sum(axis=1)
    degenerate = (valid_sum <= _EPS) & (n_sec > 0)
    scale = np.where(valid_sum > _EPS, e_sec / np.maximum(valid_sum, _EPS), 0.0)
    sec_E = sec_E * scale[:, None]
    even_share = e_sec / np.maximum(n_sec, 1).astype(np.float64)
    sec_E = np.where(degenerate[:, None] & sec_valid, even_share[:, None], sec_E)
    sec_E = sec_E.astype(np.float32)

    sec_dir_world = np.zeros((N, K, 3), dtype=np.float32)
    for i in range(K):
        valid = sec_valid[:, i]
        if valid.any():
            sec_dir_world[valid, i] = inv_local_frame_rotation(
                pre_dir[valid], dir_local[valid, i]
            )

    # mass is non-negative by construction (inv_log_transform of a real
    # number is always > 0); clip to 0 for padded/invalid slots rather than
    # leaving a spurious small positive floor from the log inverse.
    sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32)
    sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32)

    return sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid


def _physical_cond_columns(
    data: dict[str, np.ndarray], conditioning: str
) -> np.ndarray:
    """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.

    "embedding" mode zero-fills (cheap, and ConditionEncoder never reads
    these columns in that mode — so an unfilled giant.materials table can
    never crash an "embedding"-mode run). "physical" mode computes them for
    real: particle columns come from `data["mass"]`/`data["charge"]` when the
    caller already knows them directly (rollout.py, for a track descended
    from a model-predicted secondary — see giant/rollout.py's "no snapping"
    design), else derived from `data["pdg"]` via giant.particles; material
    columns always come from `data["material"]` via giant.materials, since
    material is never itself a model prediction.
    """
    from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM

    if conditioning == "embedding":
        n = len(next(iter(data.values())))
        return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32)
    if conditioning != "physical":
        raise ValueError(f"unknown conditioning mode {conditioning!r}")

    from giant.materials import material_properties_array
    from giant.particles import particle_phys_array

    if "mass" in data and "charge" in data:
        mass = np.asarray(data["mass"], dtype=np.float32)
        charge = np.asarray(data["charge"], dtype=np.float32)
    else:
        mass, charge = particle_phys_array(data["pdg"]).T

    z_eff, a_eff, density, x0, lambda_int = material_properties_array(
        data["material"]
    ).T

    return np.column_stack(
        [
            log_transform(mass),
            charge,
            z_eff,
            a_eff,
            log_transform(density),
            log_transform(x0),
            log_transform(lambda_int),
        ]
    ).astype(np.float32)


def build_cond_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    cond_normalizer: "Normalizer | None" = None,
    conditioning: str = "embedding",
) -> tuple[np.ndarray, np.ndarray]:
    """Build conditioning arrays only — no target, no post-step variables."""
    cond_cont = np.column_stack(
        [
            data["pre_pos"],
            log_transform(data["pre_E"]),
            data["pre_dir"],
            data["layer_id"].astype(np.float32),
        ]
    ).astype(np.float32)
    cond_cont = np.column_stack(
        [cond_cont, _physical_cond_columns(data, conditioning)]
    ).astype(np.float32)

    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
    mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
    cond_cat = np.column_stack([pdg_idx, mat_idx])

    if cond_normalizer is not None:
        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning)

    return cond_cont, cond_cat


def _cond_normalizer_transform(
    cond_cont: np.ndarray, cond_normalizer: "Normalizer", conditioning: str
) -> np.ndarray:
    """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed.

    Checkpoints trained before physical-property conditioning (``COND_DIM``
    8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond
    normalizer, fit before ``build_cond_features`` grew the extra physical
    columns. In "embedding" mode those columns are never read by
    ``ConditionEncoder`` (``giant/model/network.py``), so padding the missing
    entries with mean=0/std=1 is a safe no-op that keeps such checkpoints
    usable under the current, always-``COND_DIM``-wide contract. In
    "physical" mode the physical columns are load-bearing, so a mismatch
    there is a real incompatibility, not something to paper over.
    """
    mean, std = cond_normalizer.mean, cond_normalizer.std
    assert mean is not None and std is not None, "Normalizer not fitted"
    width = cond_cont.shape[-1]
    if mean.shape[-1] < width:
        if conditioning != "embedding":
            raise ValueError(
                f"cond normalizer has {mean.shape[-1]} columns, expected "
                f"{width}, and conditioning={conditioning!r} reads the "
                "physical columns directly — this checkpoint predates "
                "physical-property conditioning and can't be safely padded; "
                "retrain it under the current code."
            )
        pad = width - mean.shape[-1]
        mean = np.concatenate([mean, np.zeros(pad, dtype=mean.dtype)])
        std = np.concatenate([std, np.ones(pad, dtype=std.dtype)])
    return ((cond_cont - mean) / std).astype(np.float32)


def build_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    cond_normalizer: Normalizer | None = None,
    target_normalizer: Normalizer | None = None,
    sec_phys_normalizer: Normalizer | None = None,
    fit: bool = False,
    proc_map: dict[str, int] | None = None,
    require_secondaries: bool = False,
    conditioning: str = "embedding",
) -> tuple[
    np.ndarray,
    np.ndarray,
    np.ndarray,
    np.ndarray,
    np.ndarray,
    np.ndarray,
    Normalizer | None,
    Normalizer | None,
]:
    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) arrays.

    target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
    n_sec: (N,) integer secondary counts (target for n_sec head)
    sec_cont: (N, K_MAX, SEC_SLOT_DIM=6) continuous secondary targets
              [stick_logit, dir_local, log_mass, charge] — mass/charge are
              the secondary's real physical identity (from its ground-truth
              PDG code), a fixed regression target, not a learned/snapped one.
    proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
              never conditioning). Zeros when `proc_map` is None or the loaded
              data has no "process" column (e.g. pre-conversion parquet files).

    require_secondaries: when True, raise if any step has n_sec > 0 but the
        per-secondary list columns are absent (a mis-converted file that would
        otherwise silently zero all Stage-2 targets). Training paths set this;
        Stage-1-only callers (e.g. `giant predict`) leave it False.
    """
    from giant.constants import K_MAX

    post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
    travel_dir_local = local_frame_rotation(
        data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
    )

    energy_z = energy_simplex_encode(
        data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
    )  # (N, 2)

    target_s1 = np.column_stack(
        [
            log_transform(data["step_length"]),
            energy_z,
            post_dir_local,
            travel_dir_local,
        ]
    ).astype(np.float32)  # (N, 9)

    # Phase 2: conditioning drops n_sec and log(e_sec)
    cond_cont = np.column_stack(
        [
            data["pre_pos"],
            log_transform(data["pre_E"]),
            data["pre_dir"],
            data["layer_id"].astype(np.float32),
        ]
    ).astype(np.float32)  # (N, COND_DIM_BASE=8)
    cond_cont = np.column_stack(
        [cond_cont, _physical_cond_columns(data, conditioning)]
    ).astype(np.float32)  # (N, COND_DIM=15)

    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
    mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
    cond_cat = np.column_stack([pdg_idx, mat_idx])  # (N, 2)

    n_sec_raw = data["n_sec"].astype(
        np.int64
    )  # (N,) unclamped, for the valid-slot mask
    # Clamp the classification label to K_MAX: the head only has K_MAX+1 classes
    # (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already
    # applied to sec_cont by the loader's list padding. Without this, a rare
    # high-multiplicity step (real data goes up to ~37) hands cross_entropy
    # an out-of-range target and CUDA asserts.
    n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64)  # (N,)

    # Secondary continuous targets
    sec_E_list = data.get("sec_E_list")
    sec_dir_list = data.get("sec_dir_list")
    sec_pdg_list = data.get("sec_pdg_list")

    if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None:
        sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None]  # (N, K_MAX)
        sec_cont = encode_secondaries(
            sec_E_list,
            sec_dir_list,
            sec_valid,
            data["e_sec"],
            data["pre_dir"],
            sec_pdg_list=sec_pdg_list,
        )  # (N, K_MAX, 6)
    else:
        # Guard against silently training Stage 2 on zeroed targets: if any step
        # actually spawned secondaries (n_sec > 0, from child_track_ids) but the
        # per-secondary columns are absent, the file was never run through the
        # parent->child join (steps_to_parquet._add_secondary_attributes /
        # `dwarf convert`). Zero-filling here would collapse every secondary to
        # PDG index 0 and a constant energy fraction — a broken Stage 2 with no
        # error. Callers that only need Stage-1 (e.g. `giant predict`) keep the
        # default require_secondaries=False.
        if require_secondaries and n_sec_raw.max(initial=0) > 0:
            n_with_sec = int((n_sec_raw > 0).sum())
            raise ValueError(
                f"{n_with_sec} step(s) have secondaries (n_sec > 0) but the "
                "per-secondary columns (sec_E_list / sec_pdg_list / sec_dx_list "
                "…) are missing. This parquet was not run through the "
                "parent->child join (steps_to_parquet._add_secondary_attributes "
                "/ `dwarf convert`); training on it would silently zero all "
                "Stage-2 targets. Re-convert the file, or pass "
                "require_secondaries=False for Stage-1-only use."
            )
        N = len(n_sec)
        sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)

    if fit:
        cond_normalizer = Normalizer().fit(cond_cont)
        target_normalizer = Normalizer().fit(target_s1)

    if cond_normalizer is not None:
        cond_cont = cond_normalizer.transform(cond_cont)
    if target_normalizer is not None:
        target_s1 = target_normalizer.transform(target_s1)
    if sec_phys_normalizer is not None:
        N_, K_, _ = sec_cont.shape
        phys = sec_phys_normalizer.transform(sec_cont[:, :, 4:6].reshape(-1, 2))
        sec_cont = sec_cont.copy()
        sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)

    process = data.get("process")
    if proc_map is not None and process is not None:
        proc_idx = np.array([proc_map[str(p)] for p in process], dtype=np.int64)
    else:
        proc_idx = np.zeros(len(cond_cat), dtype=np.int64)

    return (
        cond_cont,
        cond_cat,
        target_s1,
        n_sec,
        sec_cont,
        proc_idx,
        cond_normalizer,
        target_normalizer,
    )