giant/tests/test_transforms.py

import warnings

import numpy as np
import pytest
from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX
from giant.data.transforms import (
    build_cond_features,
    build_features,
    encode_secondaries,
    energy_simplex_decode,
    energy_simplex_encode,
    inv_local_frame_rotation,
    inv_log_transform,
    local_frame_rotation,
    log_transform,
    Normalizer,
    reconstruct_post_pos,
    sorted_membership,
    travel_direction,
    _vectorized_map_lookup,
    _WelfordAccumulator,
)


def test_log_transform_invertible():
    x = np.array([0.1, 1.0, 10.0, 1000.0], dtype=np.float32)
    np.testing.assert_allclose(inv_log_transform(log_transform(x)), x, rtol=1e-5)


def test_log_transform_raises_on_input_below_negative_eps():
    """A meaningfully negative input (upstream data corruption, not float
    noise near 0) must raise instead of silently returning NaN."""
    x = np.array([1.0, -5.0], dtype=np.float32)
    with np.errstate(invalid="ignore"), pytest.raises(ValueError, match="non-finite"):
        log_transform(x)


def test_log_transform_raises_on_nan_input():
    x = np.array([1.0, np.nan], dtype=np.float32)
    with pytest.raises(ValueError, match="non-finite"):
        log_transform(x)


def test_encode_secondaries_warns_when_sec_energies_exceed_e_sec():
    """sec_E_list summing to more than e_sec (before the last slot is even
    reached) is a real upstream data mismatch — must warn instead of
    silently saturating the overflowing slot's stick-breaking logit via the
    _EPS floor. (A single slot alone exceeding what's left of the budget is
    the normal, expected "last slot takes the remainder" case and must NOT
    warn — the mismatch here is the *cumulative* sum through an earlier
    slot already exceeding e_sec.)"""
    sec_E_list = np.array([[5.0, 4.0, 1.0]], dtype=np.float32)  # sums to 10
    sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 3, 1)).astype(np.float32)
    sec_valid = np.array([[True, True, True]])
    e_sec = np.array([6.0], dtype=np.float32)  # cumsum already 9 by slot 2
    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

    with pytest.warns(UserWarning, match="sec_E_list summing to more than e_sec"):
        encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)


def test_encode_secondaries_no_warning_when_energies_are_consistent():
    sec_E_list = np.array([[3.0, 2.0]], dtype=np.float32)  # sums to 5
    sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 2, 1)).astype(np.float32)
    sec_valid = np.array([[True, True]])
    e_sec = np.array([6.0], dtype=np.float32)  # >= 5, no shortfall
    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

    with warnings.catch_warnings():
        warnings.simplefilter("error")
        encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)


def test_local_frame_rotation_noop_when_aligned():
    N = 8
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
    rng = np.random.default_rng(0)
    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
    result = local_frame_rotation(pre_dir, post_dir)
    np.testing.assert_allclose(result, post_dir, atol=1e-5)


def test_local_frame_rotation_preserves_angle():
    """Angle between pre_dir and post_dir must equal angle between ẑ and rotated."""
    rng = np.random.default_rng(1)
    N = 200
    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)

    rotated = local_frame_rotation(pre_dir, post_dir)

    cos_before = (pre_dir * post_dir).sum(axis=1)
    cos_after = rotated[:, 2]  # dot with ẑ = z-component (unit vectors)
    np.testing.assert_allclose(cos_after, cos_before, atol=1e-5)


def test_local_frame_rotation_preserves_norm():
    rng = np.random.default_rng(2)
    N = 100
    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
    result = local_frame_rotation(pre_dir, post_dir)
    np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)


def test_local_frame_rotation_rejects_near_zero_pre_dir():
    """A degenerate (near-zero-norm) pre_dir has no well-defined frame — must
    raise instead of silently falling back to an arbitrary rotation axis."""
    pre_dir = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
    post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
    with pytest.raises(ValueError, match="near-zero norm"):
        local_frame_rotation(pre_dir, post_dir)


def test_local_frame_rotation_rejects_nan_pre_dir():
    """A NaN pre_dir must raise loudly — `norm < 1e-6` is False for NaN, so
    without an explicit isfinite check this would silently poison the
    rotation (and any normalizer stats it feeds) instead of erroring."""
    pre_dir = np.array([[np.nan, 0.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32)
    post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
    with pytest.raises(ValueError, match="non-finite"):
        local_frame_rotation(pre_dir, post_dir)
    with pytest.raises(ValueError, match="non-finite"):
        inv_local_frame_rotation(pre_dir, post_dir)


def test_local_frame_rotation_antipodal_pre_dir_uses_x_axis_convention():
    """pre_dir ~ -ẑ (near-exact backscatter) is a second axis_norm~0
    degeneracy besides pre_dir ~ +ẑ; unlike the forward case, the Rodrigues
    axis-dependent terms are NOT negligible there ((1-cos_t)~2), so the x̂
    fallback is a real (if arbitrary and physically rare) convention choice
    rather than a no-op. Pin it explicitly — angle-preservation and the
    round-trip property must still hold even though the "roll" is degenerate.
    """
    pre_dir = np.array([[0.0, 0.0, -1.0]], dtype=np.float32)
    post_dir = np.array([[0.3, 0.4, 0.5]], dtype=np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)

    rotated = local_frame_rotation(pre_dir, post_dir)

    cos_before = (pre_dir * post_dir).sum(axis=1)
    cos_after = rotated[:, 2]
    np.testing.assert_allclose(cos_after, cos_before, atol=1e-5)
    np.testing.assert_allclose(np.linalg.norm(rotated, axis=1), 1.0, atol=1e-5)

    recovered = inv_local_frame_rotation(pre_dir, rotated)
    np.testing.assert_allclose(recovered, post_dir, atol=1e-5)


def test_local_frame_rotation_normalizes_non_unit_pre_dir():
    """A pre_dir with float32-drift norm (not exactly 1) must still produce the
    same result as its exactly-normalized counterpart, not a skewed frame."""
    rng = np.random.default_rng(9)
    N = 50
    pre_dir_unit = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir_unit /= np.linalg.norm(pre_dir_unit, axis=1, keepdims=True)
    post_dir = rng.standard_normal((N, 3)).astype(np.float32)
    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)

    pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(
        np.float32
    )
    expected = local_frame_rotation(pre_dir_unit, post_dir)
    result = local_frame_rotation(pre_dir_scaled, post_dir)
    np.testing.assert_allclose(result, expected, atol=1e-4)


def test_travel_direction_is_unit_norm():
    rng = np.random.default_rng(5)
    N = 50
    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
    post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
    result = travel_direction(pre_pos, post_pos)
    np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)


def test_travel_direction_matches_normalized_displacement():
    rng = np.random.default_rng(6)
    N = 50
    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
    disp = rng.standard_normal((N, 3)).astype(np.float32)
    post_pos = pre_pos + disp
    expected = disp / np.linalg.norm(disp, axis=1, keepdims=True)
    np.testing.assert_allclose(travel_direction(pre_pos, post_pos), expected, atol=1e-5)


def test_reconstruct_post_pos_straight_line():
    """When post_pos = pre_pos + L * pre_dir, travel_dir equals pre_dir, so its
    local-frame encoding is ẑ — reconstruction must recover post_pos exactly."""
    rng = np.random.default_rng(7)
    N = 20
    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
    step_length = rng.uniform(0.1, 5.0, size=N).astype(np.float32)
    post_pos = pre_pos + step_length[:, None] * pre_dir

    travel_dir_local = local_frame_rotation(
        pre_dir, travel_direction(pre_pos, post_pos)
    )
    np.testing.assert_allclose(travel_dir_local, np.tile([0, 0, 1], (N, 1)), atol=1e-4)

    reconstructed = reconstruct_post_pos(
        pre_pos, pre_dir, step_length, travel_dir_local
    )
    np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)


def test_reconstruct_post_pos_general_roundtrip():
    """Full encode (build_features-style) -> decode (cli.py predict-style) path."""
    rng = np.random.default_rng(8)
    N = 100
    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
    post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
    step_length = np.linalg.norm(post_pos - pre_pos, axis=1).astype(np.float32)

    travel_dir_local = local_frame_rotation(
        pre_dir, travel_direction(pre_pos, post_pos)
    )
    reconstructed = reconstruct_post_pos(
        pre_pos, pre_dir, step_length, travel_dir_local
    )
    np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)


def test_energy_simplex_conservation():
    """Decoding any ALR coords yields energies that sum to pre_E exactly."""
    rng = np.random.default_rng(11)
    N = 500
    z = rng.standard_normal((N, 2)).astype(np.float32) * 3.0
    pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
    edep, e_sec, post_E, delta_e = energy_simplex_decode(z, pre_E)
    np.testing.assert_allclose(edep + e_sec + post_E, pre_E, rtol=1e-5, atol=1e-4)
    np.testing.assert_allclose(delta_e, edep + e_sec, rtol=1e-5, atol=1e-4)
    assert np.all(edep >= 0) and np.all(e_sec >= 0) and np.all(post_E >= 0)


def test_energy_simplex_roundtrip():
    """Encode → decode recovers energies whose lost part already sums to delta_e."""
    rng = np.random.default_rng(12)
    N = 500000
    pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
    post_E = (pre_E * rng.uniform(0.0, 1.0, N)).astype(np.float32)
    delta_e = pre_E - post_E
    g = rng.uniform(0.0, 1.0, N).astype(np.float32)
    edep = (g * delta_e).astype(np.float32)
    e_sec = ((1.0 - g) * delta_e).astype(np.float32)

    z = energy_simplex_encode(edep, e_sec, post_E, pre_E)
    edep_r, e_sec_r, post_E_r, _ = energy_simplex_decode(z, pre_E)
    # Tolerance reflects the tiny simplex floor (~1e-5 of pre_E).
    np.testing.assert_allclose(edep_r, edep, atol=5e-3)
    np.testing.assert_allclose(e_sec_r, e_sec, atol=5e-3)
    np.testing.assert_allclose(post_E_r, post_E, atol=5e-3)


def test_energy_simplex_handles_boundary_zeros():
    """e_sec=0 (no secondaries) and post_E=0 (track end) stay finite and decode near 0."""
    pre_E = np.array([10.0, 50.0, 100.0], dtype=np.float32)
    edep = np.array([4.0, 50.0, 0.0], dtype=np.float32)
    e_sec = np.array([0.0, 0.0, 0.0], dtype=np.float32)  # no secondaries
    post_E = np.array([6.0, 0.0, 100.0], dtype=np.float32)  # row 1: track ends

    z = energy_simplex_encode(edep, e_sec, post_E, pre_E)
    assert np.all(np.isfinite(z))
    _, e_sec_r, post_E_r, _ = energy_simplex_decode(z, pre_E)
    np.testing.assert_allclose(e_sec_r, 0.0, atol=1e-2)
    assert post_E_r[1] < 1e-2  # the absorbed track decodes to ~0 post energy


def test_normalizer_roundtrip():
    rng = np.random.default_rng(3)
    X = rng.standard_normal((200, 9)).astype(np.float32)
    norm = Normalizer().fit(X)
    np.testing.assert_allclose(norm.inverse_transform(norm.transform(X)), X, atol=1e-5)


def test_normalizer_serialization():
    rng = np.random.default_rng(4)
    X = rng.standard_normal((50, 6)).astype(np.float32)
    norm = Normalizer().fit(X)
    norm2 = Normalizer.from_dict(norm.to_dict())
    assert norm2.mean is not None and norm.mean is not None
    assert norm2.std is not None and norm.std is not None
    np.testing.assert_allclose(norm2.mean, norm.mean)
    np.testing.assert_allclose(norm2.std, norm.std)


def test_build_features_clamps_n_sec_label_to_k_max():
    """A step with more secondaries than K_MAX must not overflow the
    n_sec classifier's K_MAX+1 classes (regression test: this used to hand
    cross_entropy an out-of-range target and crash CUDA training with
    'unique_by_key: failed to synchronize: cudaErrorAssert')."""
    N = 3
    raw_n_sec = np.array([0, 5, K_MAX + 20], dtype=np.int32)
    rng = np.random.default_rng(0)
    data = {
        "pdg": np.array([11, 11, 11], dtype=np.int32),
        "material": np.array(["PbWO4", "PbWO4", "PbWO4"], dtype=object),
        "pre_pos": rng.standard_normal((N, 3)).astype(np.float32),
        "pre_E": np.full(N, 10.0, dtype=np.float32),
        "pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
        "layer_id": np.zeros(N, dtype=np.int32),
        "n_sec": raw_n_sec,
        "e_sec": np.full(N, 1.0, dtype=np.float32),
        "step_length": np.full(N, 1.0, dtype=np.float32),
        "post_E": np.full(N, 9.0, dtype=np.float32),
        "edep": np.full(N, 1.0, dtype=np.float32),
        "post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
        "post_pos": rng.standard_normal((N, 3)).astype(np.float32),
    }
    pdg_map = {11: 0}
    mat_map = {"PbWO4": 0}

    _, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map)

    assert n_sec.max() <= K_MAX
    np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])


def _minimal_step_data(N: int, process: np.ndarray | None = None) -> dict:
    rng = np.random.default_rng(0)
    data = {
        "pdg": np.full(N, 11, dtype=np.int32),
        "material": np.full(N, "PbWO4", dtype=object),
        "pre_pos": rng.standard_normal((N, 3)).astype(np.float32),
        "pre_E": np.full(N, 10.0, dtype=np.float32),
        "pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
        "layer_id": np.zeros(N, dtype=np.int32),
        "n_sec": np.zeros(N, dtype=np.int32),
        "e_sec": np.full(N, 1.0, dtype=np.float32),
        "step_length": np.full(N, 1.0, dtype=np.float32),
        "post_E": np.full(N, 9.0, dtype=np.float32),
        "edep": np.full(N, 1.0, dtype=np.float32),
        "post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
        "post_pos": rng.standard_normal((N, 3)).astype(np.float32),
    }
    if process is not None:
        data["process"] = process
    return data


def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
    """Minimal build_features input with n_sec but no per-secondary list columns
    (mimics a parquet that skipped the parent->child join)."""
    data = _minimal_step_data(len(n_sec))
    data["n_sec"] = np.asarray(n_sec, dtype=np.int32)
    return data


def test_build_features_proc_idx_zero_without_proc_map():
    data = _minimal_step_data(
        3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
    )
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}

    *_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map)

    np.testing.assert_array_equal(proc_idx, [0, 0, 0])


def test_build_features_proc_idx_looks_up_proc_map():
    data = _minimal_step_data(
        3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
    )
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
    proc_map = {"compt": 0, "phot": 1, "eIoni": 2}

    *_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map)

    np.testing.assert_array_equal(proc_idx, [0, 1, 2])


def test_build_features_require_secondaries_raises_when_lists_missing():
    """A parquet with n_sec > 0 but no per-secondary list columns was never run
    through the parent->child join; require_secondaries must catch it instead of
    silently zeroing every Stage-2 target (regression: this collapsed the
    secondary species to a single PDG index during training)."""
    data = _step_data_no_sec_lists(np.array([0, 2, 1]))
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}

    with pytest.raises(ValueError, match="per-secondary columns"):
        build_features(data, pdg_map, mat_map, require_secondaries=True)


def test_build_features_require_secondaries_ok_when_no_secondaries():
    """require_secondaries only fires when secondaries actually exist; a file
    with n_sec == 0 everywhere (e.g. Stage-1-only) must still load."""
    data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}

    _, _, _, _, sec_cont, *_ = build_features(
        data, pdg_map, mat_map, require_secondaries=True
    )

    assert not sec_cont.any()


# ── physical-property conditioning ────────────────────────────────────────────


@pytest.fixture
def fake_material_props(monkeypatch):
    """Inject a fully-populated fake materials table for "physical" mode
    tests, independent of when the real giant/materials.py table is filled
    in by the user (see giant.materials.MaterialPropertiesNotFilledError)."""
    import giant.materials as gm

    fake = {
        "PbWO4": gm.MaterialProperties(
            z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
        )
    }
    monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
    return fake


def test_build_features_embedding_mode_zero_fills_physical_columns():
    data = _minimal_step_data(3)
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}

    cond_cont, *_ = build_features(
        data,
        pdg_map,
        mat_map,
        particle_conditioning="embedding",
        material_conditioning="embedding",
    )

    assert cond_cont.shape[1] == COND_DIM
    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)


def test_build_features_physical_mode_shape_and_values(fake_material_props):
    from giant.particles import particle_mass_charge

    data = _minimal_step_data(3)
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}

    cond_cont, *_ = build_features(
        data,
        pdg_map,
        mat_map,
        particle_conditioning="physical",
        material_conditioning="physical",
    )

    assert cond_cont.shape[1] == COND_DIM
    mass, charge = particle_mass_charge(11)
    expected_log_mass = log_transform(np.array([mass]))[0]
    np.testing.assert_allclose(
        cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5
    )
    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], charge)
    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 2], 75.6)  # z_eff
    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 3], 205.3)  # a_eff


def test_build_features_physical_mode_unfilled_material_raises():
    """G4_LYSO is the one material giant/materials.py still ships unfilled
    (not a stock Geant4 NIST material) — must fail loudly, not silently."""
    from giant.materials import MaterialPropertiesNotFilledError

    data = _minimal_step_data(2)
    data["material"] = np.full(2, "G4_LYSO", dtype=object)
    pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0}

    with pytest.raises(MaterialPropertiesNotFilledError):
        build_features(
            data,
            pdg_map,
            mat_map,
            particle_conditioning="physical",
            material_conditioning="physical",
        )


def test_build_cond_features_mass_charge_override(fake_material_props):
    """rollout.py's secondaries carry their own predicted mass/charge — when
    present in `data`, these bypass the pdg-based lookup entirely (the "no
    snapping" design: a track's own future conditioning must use its actual
    predicted physical identity, not a value re-derived from a PDG code)."""
    data = _minimal_step_data(2)
    data["mass"] = np.array([123.0, 456.0], dtype=np.float32)
    data["charge"] = np.array([2.0, -2.0], dtype=np.float32)
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}

    cond_cont, _ = build_cond_features(
        data,
        pdg_map,
        mat_map,
        particle_conditioning="physical",
        material_conditioning="physical",
    )

    np.testing.assert_allclose(
        cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
    )
    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0])


def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode():
    """A pre-physical-conditioning checkpoint's cond normalizer is COND_DIM_BASE
    (8) wide, fit before build_cond_features grew the extra physical columns.
    In "embedding" mode those columns are never read downstream, so a legacy
    normalizer should be usable as-is (padded, not rejected)."""
    data = _minimal_step_data(3)
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
    legacy_norm = Normalizer()
    legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
    legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)

    cond_cont, _ = build_cond_features(
        data,
        pdg_map,
        mat_map,
        cond_normalizer=legacy_norm,
        particle_conditioning="embedding",
        material_conditioning="embedding",
    )

    assert cond_cont.shape[-1] == COND_DIM
    # padded physical columns are zero-filled pre-normalization and
    # mean=0/std=1 post-normalization, so they should come out as zero
    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)


def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
    fake_material_props,
):
    """Unlike "embedding" mode, "physical" mode actually reads the physical
    columns, so a legacy 8-wide normalizer can't be silently padded — that
    would silently feed the network un-normalized physical properties."""
    data = _minimal_step_data(3)
    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
    legacy_norm = Normalizer()
    legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
    legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)

    with pytest.raises(ValueError, match="predates physical-property conditioning"):
        build_cond_features(
            data,
            pdg_map,
            mat_map,
            cond_normalizer=legacy_norm,
            particle_conditioning="physical",
            material_conditioning="physical",
        )


# ── sorted_membership / _vectorized_map_lookup ──────────────────────────────


def test_sorted_membership_matches_np_isin():
    rng = np.random.default_rng(0)
    sorted_arr = np.unique(rng.integers(0, 10_000, size=500))
    values = rng.integers(-100, 10_100, size=2_000)  # some in, some out of range
    # values deliberately not sorted
    rng.shuffle(values)

    result = sorted_membership(values, sorted_arr)
    expected = np.isin(values, sorted_arr)
    np.testing.assert_array_equal(result, expected)


def test_sorted_membership_empty_sorted_arr():
    values = np.array([1, 2, 3])
    sorted_arr = np.array([], dtype=np.int64)
    result = sorted_membership(values, sorted_arr)
    np.testing.assert_array_equal(result, np.zeros(3, dtype=bool))


def test_vectorized_map_lookup_matches_dict_comprehension_int_keys():
    rng = np.random.default_rng(1)
    keys = np.unique(rng.integers(-1000, 1000, size=200))
    mapping = {int(k): i for i, k in enumerate(keys)}
    values = rng.choice(keys, size=500)

    result = _vectorized_map_lookup(values, mapping)
    expected = np.array([mapping[int(v)] for v in values], dtype=np.int64)
    np.testing.assert_array_equal(result, expected)


def test_vectorized_map_lookup_matches_dict_comprehension_str_keys():
    mapping = {"PbWO4": 0, "G4_AIR": 1, "G4_Fe": 2}
    values = np.array(["G4_Fe", "PbWO4", "G4_AIR", "PbWO4"], dtype=object)

    result = _vectorized_map_lookup(values, mapping)
    expected = np.array([mapping[str(v)] for v in values], dtype=np.int64)
    np.testing.assert_array_equal(result, expected)


def test_vectorized_map_lookup_raises_keyerror_on_missing_value():
    mapping = {1: 0, 2: 1}
    values = np.array([1, 2, 3])
    with pytest.raises(KeyError):
        _vectorized_map_lookup(values, mapping)


def test_vectorized_map_lookup_strict_false_dummy_indexes_unmapped_values():
    """strict=False must leave found values untouched and only dummy-index
    (0) the unmapped ones — never raise, and never disturb a value that IS
    in the mapping (e.g. one that happens to map to a nonzero index)."""
    mapping = {1: 5, 2: 7}
    values = np.array([1, 99, 2, 100])
    result = _vectorized_map_lookup(values, mapping, strict=False)
    np.testing.assert_array_equal(result, [5, 0, 7, 0])


def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_material():
    """conditioning="physical" must not KeyError on a pdg/material outside
    the training-dataset vocab (mat_map/pdg_map) — that's the entire point
    of the mode (see giant.rollout's known_pdg gate for the paired fix).
    "embedding" mode must still raise, since cond_cat IS the conditioning
    signal there. Note this is specifically about the dataset-scoped
    vocab index, not giant.materials' physical-properties table — a
    material must still be a real, known Geant4 material (e.g. "G4_Pb",
    just not one *this* mat_map happened to include) for "physical" mode
    to derive its Z_eff/A_eff/density/X0/λ_int; a genuinely unknown
    material name correctly still raises via giant.materials, same as the
    documented G4_LYSO precedent — that's a separate, intentional guard."""
    pdg_map = {11: 0, 22: 1}
    mat_map = {"G4_AIR": 0}
    data = {
        "pre_pos": np.zeros((1, 3), dtype=np.float32),
        "pre_E": np.array([10.0], dtype=np.float32),
        "pre_dir": np.array([[0.0, 0.0, 1.0]], dtype=np.float32),
        "layer_id": np.array([0], dtype=np.int32),
        "pdg": np.array([13], dtype=np.int64),  # not in pdg_map
        "material": np.array(["G4_Pb"], dtype=object),  # not in mat_map
        "mass": np.array([105.7], dtype=np.float32),
        "charge": np.array([-1.0], dtype=np.float32),
    }

    cond_cont, cond_cat = build_cond_features(
        data,
        pdg_map,
        mat_map,
        particle_conditioning="physical",
        material_conditioning="physical",
    )
    assert cond_cont.shape[-1] == COND_DIM
    np.testing.assert_array_equal(cond_cat, [[0, 0]])  # dummy indices, no raise

    with pytest.raises(KeyError):
        build_cond_features(
            data,
            pdg_map,
            mat_map,
            particle_conditioning="embedding",
            material_conditioning="embedding",
        )


# ── _WelfordAccumulator ──────────────────────────────────────────────────────


def test_welford_accumulator_matches_direct_mean_std_over_many_chunks():
    rng = np.random.default_rng(5)
    F = 4
    chunks = [rng.standard_normal((rng.integers(1, 50), F)) * 10 + 3 for _ in range(20)]
    full = np.concatenate(chunks, axis=0)

    acc = _WelfordAccumulator(F)
    for chunk in chunks:
        acc.update(chunk)
    norm = acc.to_normalizer()

    assert norm.mean is not None and norm.std is not None
    np.testing.assert_allclose(norm.mean, full.mean(axis=0), rtol=1e-5, atol=1e-5)
    np.testing.assert_allclose(norm.std, full.std(axis=0), rtol=1e-5, atol=1e-5)
    assert acc.n == full.shape[0]


def test_welford_accumulator_single_chunk():
    rng = np.random.default_rng(6)
    X = rng.standard_normal((100, 3)) * 5 - 2

    acc = _WelfordAccumulator(3)
    acc.update(X)
    norm = acc.to_normalizer()

    assert norm.mean is not None and norm.std is not None
    np.testing.assert_allclose(norm.mean, X.mean(axis=0), rtol=1e-5)
    np.testing.assert_allclose(norm.std, X.std(axis=0), rtol=1e-5)


def test_welford_accumulator_matches_naive_running_mean_reference():
    """The chunk-local-mean + Chan-merge formula must agree with the naive
    textbook streaming update (subtract the *running* mean before and after
    updating it) that it replaces, within float64 rounding tolerance."""
    rng = np.random.default_rng(7)
    F = 3
    chunks = [rng.standard_normal((rng.integers(1, 40), F)) for _ in range(15)]

    def naive_update(mean, M2, n, X):
        X = np.asarray(X, dtype=np.float64)
        B = X.shape[0]
        new_n = n + B
        delta = X - mean
        mean = mean + delta.sum(0) / new_n
        delta2 = X - mean
        M2 = M2 + (delta * delta2).sum(0)
        return mean, M2, new_n

    naive_mean = np.zeros(F)
    naive_M2 = np.zeros(F)
    naive_n = 0
    for chunk in chunks:
        naive_mean, naive_M2, naive_n = naive_update(
            naive_mean, naive_M2, naive_n, chunk
        )

    acc = _WelfordAccumulator(F)
    for chunk in chunks:
        acc.update(chunk)

    assert acc.n == naive_n
    np.testing.assert_allclose(acc._mean, naive_mean, rtol=1e-9, atol=1e-9)
    np.testing.assert_allclose(acc._M2, naive_M2, rtol=1e-9, atol=1e-9)