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


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,
) -> 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, 4):
      slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z]

    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.

    sec_pdg_idx (integer) is not processed here — kept separate so the loss
    function can look up the embedding table at training time.
    """
    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]
            )

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


def decode_secondaries(
    sec_cont: np.ndarray,
    sec_pdg_pred: np.ndarray,
    n_sec: np.ndarray,
    e_sec: np.ndarray,
    pre_dir: np.ndarray,
    pdg_map_inv: dict[int, int],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Inverse of encode_secondaries: continuous targets → physical secondary attrs.

    sec_cont: (N, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
    sec_pdg_pred: (N, K_MAX) integer PDG indices (from nearest-neighbor snap)
    n_sec: (N,) integer secondary counts
    e_sec: (N,) total secondary energy budget [MeV]
    pre_dir: (N, 3) pre-step world-frame direction
    pdg_map_inv: maps model index → PDG code

    Returns (sec_E, sec_dir_world, sec_pdg_code, 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.
    """
    N, K, _ = sec_cont.shape
    stick_logits = sec_cont[:, :, 0]  # (N, K)
    dir_local = sec_cont[:, :, 1:].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]
            )

    sec_pdg_code = np.array(
        [
            [pdg_map_inv.get(int(sec_pdg_pred[n, i]), 0) for i in range(K)]
            for n in range(N)
        ],
        dtype=np.int32,
    )

    return sec_E, sec_dir_world, sec_pdg_code, sec_valid


def build_cond_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    cond_normalizer: "Normalizer | None" = None,
) -> 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)

    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)

    return cond_cont, cond_cat


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,
    fit: bool = False,
    require_secondaries: bool = False,
) -> 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, sec_pdg_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, 4) continuous secondary targets [stick_logit, dir_local]
    sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding
                 targets in the training loop

    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=8)

    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/sec_pdg_idx 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"]
        )  # (N, K_MAX, 4)
        # Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int),
        # which is never a real PDG code, so `.get(..., 0)` naturally maps
        # both real unknown codes and padding to the same masked-out index.
        sec_pdg_idx = np.vectorize(lambda p: pdg_map.get(int(p), 0))(
            sec_pdg_list
        ).astype(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, 4), dtype=np.float32)
        sec_pdg_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)
    if target_normalizer is not None:
        target_s1 = target_normalizer.transform(target_s1)

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