giant/giant/data/transforms.py

import warnings
from typing import NamedTuple

import numpy as np

from giant.cond_layout import CondLayout
from giant.constants import K_MAX

_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:
    x = np.asarray(x, dtype=np.float32)
    y = np.log(x + eps)
    if not np.all(np.isfinite(y)):
        bad = int(np.sum(~np.isfinite(y)))
        raise ValueError(
            f"log_transform: {bad} value(s) produced non-finite output (input "
            f"< -eps={eps:g}, or already NaN/Inf); every quantity this is "
            "applied to should be non-negative, so this indicates upstream "
            "data corruption rather than expected float noise."
        )
    return y


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)
    # axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
    # pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
    # ẑ) vanishes at both. The "choice is irrelevant" claim below only holds
    # at +ẑ, where sin_t~0 AND (1-cos_t)~0 so every axis-dependent Rodrigues
    # term vanishes. At -ẑ, sin_t~0 but (1-cos_t)~2 — not negligible — so
    # snapping to a fixed x̂ there is a genuine (if physically rare)
    # modeling choice, not a no-op: it picks one representative out of an
    # inherently ambiguous family of 180°-about-any-transverse-axis
    # rotations (no single-valued frame convention can be continuous through
    # this antipode — same obstruction as a sphere's tangent frame having no
    # continuous choice at a pole). x̂ is still fine to use — it's a fixed,
    # self-consistent convention that `local_frame_rotation`/
    # `inv_local_frame_rotation` (same threshold) round-trip correctly
    # through — but steps whose pre_dir falls in this tiny near-backscatter
    # cone get a discontinuous "roll" relative to their non-degenerate
    # neighbors, injecting a small amount of label noise there.
    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).

    NaN/Inf rows are also raised on explicitly: `norm < 1e-6` is False for a
    NaN norm, so without this check a non-finite row would silently pass
    through and poison everything downstream (e.g. the persisted normalizer
    stats in `setup_cache`, if the row is swept into a Welford accumulator).
    """
    pre_dir = np.asarray(pre_dir, dtype=np.float32)
    if not np.all(np.isfinite(pre_dir)):
        bad = int(np.sum(~np.all(np.isfinite(pre_dir), axis=1)))
        raise ValueError(
            f"pre_dir has {bad} row(s) with non-finite (NaN/Inf) components; "
            "local/inv_local_frame_rotation require a well-defined incoming "
            "direction for every row."
        )
    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 (Chan/Golub/LeVeque 1979 parallel algorithm).

    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:
        # Computes the chunk's own local mean/M2 (two passes over X, no
        # reference to the running mean) and merges it into the running
        # totals with the O(F) Chan/Golub/LeVeque combination formula.
        # Equivalent to the textbook single-pass streaming update (which
        # instead re-derives two full (B, F) arrays from the running mean,
        # before and after updating it) but ~40% cheaper here since it
        # avoids one of those (B, F) passes and its temporary array.
        X = np.asarray(X, dtype=np.float64)
        B = X.shape[0]
        mean_b = X.mean(0)
        diff = X - mean_b
        M2_b = np.einsum("ij,ij->j", diff, diff)

        new_n = self.n + B
        delta = mean_b - self._mean
        self._mean += delta * (B / new_n)
        self._M2 += M2_b + delta * delta * (self.n * B / new_n)
        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 sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
    """Boolean membership of `values` (any order) in `sorted_arr` (ascending, unique).

    Equivalent to `np.isin(values, sorted_arr)`, but `np.isin`'s default path
    sorts both inputs on every call — costly when `sorted_arr` is a large,
    already-sorted array (e.g. all train-split event ids) reused across many
    chunks. This does one `searchsorted` per call instead. `values` need not
    be sorted; `sorted_arr` must be ascending and duplicate-free.
    """
    values = np.asarray(values)
    if sorted_arr.size == 0:
        return np.zeros(values.shape, dtype=bool)
    idx = np.searchsorted(sorted_arr, values)
    idx = np.clip(idx, 0, len(sorted_arr) - 1)
    return sorted_arr[idx] == values


def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0) -> np.ndarray:
    """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.

    Replaces a per-element Python dict lookup with one `searchsorted` call.
    Raises `KeyError` if any value in `values` isn't a key of `mapping`,
    matching the dict-comprehension it replaces (never silently misassigns)
    — unless `strict=False`, in which case unmapped values get `default`
    instead. Only pass `strict=False` where the caller has independently
    verified the resulting index is either never actually read (e.g.
    `build_cond_features` under `conditioning="physical"`, where
    `ConditionEncoder` ignores `cond_cat` entirely) or where `default` is a
    deliberate fallback class (e.g. a top-N map's "other" index for a raw
    value outside the training vocab). It exists so a rollout can be seeded
    with a species/material outside the training vocab without a spurious
    `KeyError`, which is the entire point of physical-property conditioning.
    """
    keys = np.asarray(list(mapping.keys()))
    vals = np.asarray(list(mapping.values()), dtype=np.int64)
    order = np.argsort(keys, kind="stable")
    keys_sorted, vals_sorted = keys[order], vals[order]
    values = np.asarray(values)
    pos = np.searchsorted(keys_sorted, values)
    pos = np.clip(pos, 0, len(keys_sorted) - 1)
    found = keys_sorted[pos] == values
    if not found.all():
        if not strict:
            out = np.full(values.shape, default, dtype=np.int64)
            out[found] = vals_sorted[pos[found]]
            return out
        missing = np.unique(values[~found])
        raise KeyError(f"value(s) not in mapping: {missing[:10].tolist()}")
    return vals_sorted[pos]


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,
    phys_only: bool = False,
) -> 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.

    `phys_only=True` skips the stick-breaking and direction-rotation blocks
    (zero-filling them instead) and computes only log_mass/charge — for
    callers (normalizer fitting) that discard the other four columns anyway,
    so computing them would be wasted work repeated over the whole dataset.
    """
    N, K = sec_E_list.shape
    e_sec = np.asarray(e_sec, dtype=np.float64)

    if phys_only:
        stick_logits = np.zeros((N, K), dtype=np.float32)
        dir_local = np.zeros((N, K, 3), dtype=np.float32)
    else:
        cumsum = np.cumsum(sec_E_list.astype(np.float64), axis=1)
        stick_logits = np.zeros((N, K), dtype=np.float32)
        # A valid slot whose cumulative secondary energy so far exceeds
        # e_sec by more than float noise means sec_E_list sums to more than
        # e_sec — a real upstream data mismatch, not something to paper
        # over. Flagged once after the loop rather than let `remaining`'s
        # np.maximum(..., _EPS) floor silently absorb it by saturating that
        # slot's stick-breaking logit with no signal that anything was off.
        _SHORTFALL_TOL = 1e-3
        shortfall_flagged = np.zeros(N, dtype=bool)
        for i in range(K):
            if i == 0:
                remaining_raw = e_sec
            else:
                remaining_raw = e_sec - cumsum[:, i - 1]
            shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
            remaining = np.maximum(remaining_raw, _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)

        if shortfall_flagged.any():
            n = int(shortfall_flagged.sum())
            warnings.warn(
                f"encode_secondaries: {n}/{N} row(s) have sec_E_list summing "
                "to more than e_sec (beyond float noise) — the overflowing "
                "slot(s)' stick-breaking logit was saturated instead of "
                "reflecting a real fraction; check upstream secondary "
                "energy accounting for these rows.",
                stacklevel=2,
            )

        # 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 encode_secondary_type_idx(sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict) -> np.ndarray:
    """Per-secondary-slot class index into `class_map` — (N, K_MAX) int64.

    `class_map` is either a top-N-plus-other map's `class_map`
    (`stage2_model.particle_type.target = "onehot"`, see
    `giant.data.loader.build_pdg_topn_map_from_files`) or the dense `pdg_map`
    (`target = "embedding"`). Not used at all for `target = "physical"` —
    that target keeps using `encode_secondaries`'s (log_mass, charge)
    columns unchanged.

    Padding slots get index 0 (their looked-up value is discarded downstream
    by the `sec_valid`/`n_sec` mask regardless, so any in-vocabulary dummy
    code works). A *real, valid* secondary whose code is missing from
    `class_map` raises `KeyError` (`strict=True`) rather than silently
    misassigning — for `target="onehot"` this should never actually
    trigger, since `build_pdg_topn_map_from_files` pools both primary and
    secondary occurrences precisely so every secondary species seen in
    these files has a key (in "other" at worst); for `target="embedding"`
    (which reuses the dense, primary-only `pdg_map`) it's a real signal
    that a secondary-only species exists with no primary-role counterpart.
    """
    N, K = sec_pdg_list.shape
    # An arbitrary already-present key works as the padding-slot dummy code
    # (unlike encode_secondaries' physics-derived phys lookup, this is an
    # index into class_map's own vocabulary, so a fixed sentinel like 22
    # isn't guaranteed to be a key — an arbitrary present one always is).
    dummy = next(iter(class_map))
    safe_pdg = np.where(sec_valid, sec_pdg_list, dummy)
    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(N, K)
    return np.where(sec_valid, idx, 0).astype(np.int64)


def decode_secondary_cont(
    sec_cont: np.ndarray,
    n_sec: np.ndarray,
    e_sec: np.ndarray,
    pre_dir: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Continuous-only half of `decode_secondaries`'s inverse: the
    stick-breaking energy split and local->world direction — generator/
    `particle_type.target`-independent, since every target (`"physical"`,
    `"onehot"`, `"embedding"`) shares the same `CONT_SLOT_DIM`-wide
    (stick_logit, dir) prefix and differs only
    in what follows it. `decode_secondaries` (target="physical") is the
    original all-in-one form built on top of this; `target` in `("onehot",
    "embedding")` decodes their type slice separately via
    `giant.particles.decode_topn_class`/`decode_embedding_nearest` and calls
    this directly instead — see `giant/rollout.py`.

    sec_cont: (N, K, >=CONT_SLOT_DIM) — only columns `[:, :, :CONT_SLOT_DIM]`
        (stick_logit, local dir) are read; a caller may pass its full
        per-slot tensor (continuous + type) unsliced.
    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_valid), shapes (N, K), (N, K, 3),
    (N, K). The valid slots' energies (`sec_E[sec_valid]`, per row) always
    sum to exactly `e_sec` — see the rescaling below.
    """
    N, K = sec_cont.shape[0], sec_cont.shape[1]
    stick_logits = sec_cont[:, :, 0]  # (N, K)
    dir_local = sec_cont[:, :, 1:4].copy()  # (N, K, 3)

    # 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])

    return sec_E, sec_dir_world, sec_valid


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.

    `particle_type.target = "physical"` only (the type slice is a raw
    (log_mass, charge) regression target folded straight into `sec_cont`) —
    `"onehot"`/`"embedding"` decode through `decode_secondary_cont` +
    `giant.particles.decode_topn_class`/`decode_embedding_nearest` instead,
    since their type slice isn't (log_mass, charge) at all. See
    `decode_secondary_cont`'s docstring for why the two share the energy/
    direction logic below.

    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). 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)

    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont, n_sec, e_sec, pre_dir)

    log_mass = sec_cont[:, :, 4]  # (N, K)
    charge = sec_cont[:, :, 5]  # (N, K)

    # 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], layout: CondLayout) -> np.ndarray:
    """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.

    The particle and material blocks are gated independently and may mix
    freely — e.g. material `physical` with particle `embedding` — so e.g.
    `particle_type="embedding"` + `material_type="physical"` zero-fills only
    the particle columns and computes the material ones for real.

    "embedding"/"onehot" zero-fill their block (cheap, and ConditionEncoder
    never reads these columns in either mode — so an unfilled
    giant.materials table can never crash an "embedding"/"onehot"-mode run).
    "physical" computes it 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

    n = len(next(iter(data.values())))

    if layout.particle_type == "physical":
        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
        particle_cols = np.column_stack([log_transform(mass), charge])
    else:
        particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)

    if layout.material_type == "physical":
        from giant.materials import material_properties_array

        z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T
        material_cols = np.column_stack(
            [
                z_eff,
                a_eff,
                log_transform(density),
                log_transform(x0),
                log_transform(lambda_int),
            ]
        )
    else:
        material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)

    return np.column_stack([particle_cols, material_cols]).astype(np.float32)


def _build_cond_arrays(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    layout: CondLayout,
    pdg_topn_map: dict[int, int] | None,
    mat_topn_map: dict[str, int] | None,
) -> tuple[np.ndarray, np.ndarray]:
    """The un-normalized `(cond_cont, cond_cat)` pair, in `layout`'s column order.

    Both `build_cond_features` and `build_features` go through here, so the
    column order — and everything that depends on it — is stated once. See
    `giant.cond_layout.CondLayout` for the layout itself.
    """
    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, layout)]).astype(
        np.float32
    )  # (N, COND_DIM=15)

    # In "physical" mode cond_cat's first two columns are only a
    # reporting/router convenience — ConditionEncoder never reads them
    # (giant/model/encoders.py) — so a species/material outside the training
    # vocab (the whole point of physical-property conditioning) gets a dummy
    # index instead of raising. In "embedding" mode those columns ARE the
    # conditioning signal, so an unmapped value must still raise loudly
    # rather than silently misassign. In "onehot" mode they again go unread
    # (the topN columns below are the real signal), so they're as permissive
    # as "physical". Each axis's strictness is independent.
    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=layout.particle_type == "embedding")
    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=layout.material_type == "embedding")
    # Which extra columns exist is the layout's call, not "did the caller
    # happen to pass a map" — that's what used to let the producer and
    # ConditionEncoder disagree. A map for a non-"onehot" axis is unused.
    cat_cols = [pdg_idx, mat_idx]
    if layout.particle_topn_col is not None:
        if pdg_topn_map is None:
            raise ValueError("conditioning.particle.type='onehot' needs pdg_topn_map")
        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
    if layout.material_topn_col is not None:
        if mat_topn_map is None:
            raise ValueError("conditioning.material.type='onehot' needs mat_topn_map")
        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
    cond_cat = np.column_stack(cat_cols)  # (N, layout.cat_dim)

    return cond_cont, cond_cat


def build_cond_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    cond_normalizer: "Normalizer | None" = None,
    particle_conditioning: str = "embedding",
    material_conditioning: str = "embedding",
    pdg_topn_map: dict[int, int] | None = None,
    mat_topn_map: dict[str, int] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Build conditioning arrays only — no target, no post-step variables.

    `particle_conditioning`/`material_conditioning` are independent —
    e.g. `particle_conditioning="embedding"` +
    `material_conditioning="physical"` is a valid mix.

    `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
    `giant.data.loader.build_topn_map_from_files`) supply the extra `cond_cat`
    columns read by `ConditionEncoder`'s `"onehot"` mode, and are required
    whenever the corresponding axis is `"onehot"`. See
    `giant.cond_layout.CondLayout` for which columns exist where.
    """
    layout = CondLayout.from_types(particle_conditioning, material_conditioning)
    cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map)

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

    return cond_cont, cond_cat


def _cond_normalizer_transform(
    cond_cont: np.ndarray,
    cond_normalizer: "Normalizer",
    layout: CondLayout,
) -> 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. When NEITHER axis is "physical" those columns are never read by
    ``ConditionEncoder`` (``giant/model/encoders.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. If EITHER
    axis is "physical" its 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:
        physical_load_bearing = "physical" in (
            layout.particle_type,
            layout.material_type,
        )
        if physical_load_bearing:
            raise ValueError(
                f"cond normalizer has {mean.shape[-1]} columns, expected "
                f"{width}, and particle_conditioning={layout.particle_type!r}/"
                f"material_conditioning={layout.material_type!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)


class StepFeatures(NamedTuple):
    """Output of `build_features`. Field order is load-bearing for existing
    positional unpacking (tests, `StreamingStepsDataset`) — append only,
    never insert or reorder.

    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.
              Always computed the same way regardless of
              `stage2_model.particle_type.target` — only actually used
              downstream under `target = "physical"`.
    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).
    sec_type_idx: (N, K_MAX) integer secondary class index into
              `sec_type_class_map`, for `stage2_model.particle_type.target`
              in `("onehot", "embedding")` — see `encode_secondary_type_idx`.
              Zero-filled (and unused) when `sec_type_class_map` is None
              (i.e. `target = "physical"`).
    """

    cond_cont: np.ndarray
    cond_cat: np.ndarray
    target_s1: np.ndarray
    n_sec: np.ndarray
    sec_cont: np.ndarray
    proc_idx: np.ndarray
    sec_type_idx: np.ndarray
    cond_normalizer: Normalizer | None
    target_normalizer: Normalizer | None


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,
    particle_conditioning: str = "embedding",
    material_conditioning: str = "embedding",
    sec_phys_only: bool = False,
    pdg_topn_map: dict[int, int] | None = None,
    mat_topn_map: dict[str, int] | None = None,
    sec_type_class_map: dict | None = None,
    k_max: int = K_MAX,
) -> StepFeatures:
    """Assemble a `StepFeatures` of (cond_cont, cond_cat, target_s1, n_sec,
    sec_cont, proc_idx, sec_type_idx, cond_normalizer, target_normalizer) —
    see `StepFeatures` for field meanings.

    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.

    sec_phys_only: passed straight through to `encode_secondaries` — skips
        the stick-breaking/direction-rotation blocks of `sec_cont` (zero-filled
        instead) for callers (normalizer fitting) that only read
        `sec_cont[:, :, 4:6]` and would otherwise discard that work.

    pdg_topn_map/mat_topn_map: source of the extra `cond_cat` columns for
        `ConditionEncoder`'s `"onehot"` mode — see `build_cond_features`.

    sec_type_class_map: the map `sec_type_idx` is looked up against — a
        top-N-plus-other map's `class_map` for `target = "onehot"`, or the
        dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself).
        `None` for `target = "physical"`.

    k_max: should match `stage2_model.k_max` —
        overridden internally by `data["sec_E_list"]`'s own padded width when
        present (the loader already padded it to some k_max; that width is
        authoritative), so this only actually matters when secondary list
        columns are absent (Stage-1-only reads, or a pre-secondary-join
        file), where it sets `sec_cont`/`sec_type_idx`'s zero-filled width.
    """

    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)
    layout = CondLayout.from_types(particle_conditioning, material_conditioning)
    cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map)

    n_sec_raw = data["n_sec"].astype(np.int64)  # (N,) unclamped, for the valid-slot mask

    # 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:
        # The loader already padded sec_*_list to some k_max (see
        # giant.data.loader.iter_file_chunks); that padded width is
        # authoritative over whatever this call happened to pass in, so the
        # two can never drift apart.
        k_max = sec_E_list.shape[1]

    # 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,)

    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,
            phys_only=sec_phys_only,
        )  # (N, k_max, 6)
        sec_type_idx = (
            encode_secondary_type_idx(sec_pdg_list, sec_valid, sec_type_class_map)
            if sec_type_class_map is not None
            else np.zeros((len(n_sec), k_max), dtype=np.int64)
        )
    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)
        sec_type_idx = np.zeros((N, k_max), dtype=np.int64)

    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, cond_normalizer, layout)
    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 = _vectorized_map_lookup(process, proc_map)
    else:
        proc_idx = np.zeros(len(cond_cat), dtype=np.int64)

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