giant/tests/test_setup_cache.py

import json
import os
import stat

import numpy as np
import pandas as pd
import pytest

from giant.data import setup_cache
from giant.data.loader import TopNMap
from giant.data.setup_cache import NormalizerEntry, SetupCache
from giant.data.transforms import Normalizer


def _touch_parquet(path, n=1):
    path.parent.mkdir(parents=True, exist_ok=True)
    pd.DataFrame({"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}).to_parquet(path)
    return path


def _normalizer(width=3):
    norm = Normalizer()
    norm.mean = np.zeros(width, dtype=np.float32)
    norm.std = np.ones(width, dtype=np.float32)
    return norm


def _entry(n_train_steps=100, sample=None):
    sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) if sample is None else sample
    return NormalizerEntry(_normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample)


# ── sidecar_path ─────────────────────────────────────────────────────────


def test_sidecar_path_single_file(tmp_path):
    f = tmp_path / "shard.parquet"
    assert setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"


def test_sidecar_path_directory(tmp_path):
    d = tmp_path / "pbwo4"
    assert setup_cache.sidecar_path(d) == tmp_path / "pbwo4.giant_train_cache.json"


def test_sidecar_path_manifest(tmp_path):
    m = tmp_path / "pools" / "full.manifest"
    assert setup_cache.sidecar_path(m) == tmp_path / "pools" / "full.manifest.giant_train_cache.json"


# ── fingerprint_files ────────────────────────────────────────────────────


def test_fingerprint_files_order_preserving(tmp_path):
    a = _touch_parquet(tmp_path / "a.parquet")
    b = _touch_parquet(tmp_path / "b.parquet")

    forward = setup_cache.fingerprint_files([a, b])
    backward = setup_cache.fingerprint_files([b, a])

    assert forward[0][0] == str(a.resolve())
    assert forward[1][0] == str(b.resolve())
    assert backward[0][0] == str(b.resolve())
    assert forward != backward


# ── save / load round trip ──────────────────────────────────────────────


def test_save_load_round_trip(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]

    cache = SetupCache.empty(files)
    cache.vocab = ({11: 0, 22: 1}, {"G4_AIR": 0})
    cache.event_index = (np.array([1, 2, 3]), np.array([10, 20, 30]))
    cache.proc_maps[4] = {"eIoni": 0, "phot": 1}
    cache.normalizers["valfrac=0.1_seed=0_cond=physical"] = _entry()

    setup_cache.save(data, files, cache)
    loaded = setup_cache.load(data, files)

    assert loaded is not None
    assert loaded.vocab == ({11: 0, 22: 1}, {"G4_AIR": 0})
    assert loaded.event_index is not None
    np.testing.assert_array_equal(loaded.event_index[0], [1, 2, 3])
    np.testing.assert_array_equal(loaded.event_index[1], [10, 20, 30])
    assert loaded.proc_maps == {4: {"eIoni": 0, "phot": 1}}
    entry = loaded.normalizers["valfrac=0.1_seed=0_cond=physical"]
    assert entry.cond_norm.mean is not None
    np.testing.assert_allclose(entry.cond_norm.mean, np.zeros(3, dtype=np.float32))
    assert entry.n_train_steps == 100
    np.testing.assert_allclose(entry.energy_quantiles, [1.0, 2.0, 3.0])


def test_save_load_round_trip_topn_maps(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]

    cache = SetupCache.empty(files)
    cache.topn_maps[setup_cache.topn_key("pdg", 3)] = TopNMap(
        class_map={22: 0, 11: 1, 2212: 2}, other_members={2212: 5}
    )
    cache.topn_maps[setup_cache.topn_key("material", 2)] = TopNMap(
        class_map={"G4_AIR": 0, "PbWO4": 1}, other_members={}
    )

    setup_cache.save(data, files, cache)
    loaded = setup_cache.load(data, files)

    assert loaded is not None
    pdg_m = loaded.topn_maps[setup_cache.topn_key("pdg", 3)]
    assert pdg_m.class_map == {22: 0, 11: 1, 2212: 2}
    assert pdg_m.other_members == {2212: 5}
    # key type is int (matches pdg_map's own key type), not str
    assert all(isinstance(k, int) for k in pdg_m.class_map)

    mat_m = loaded.topn_maps[setup_cache.topn_key("material", 2)]
    assert mat_m.class_map == {"G4_AIR": 0, "PbWO4": 1}


def test_topn_key_unknown_axis_raises():
    with pytest.raises(ValueError, match="unknown top-N map axis"):
        setup_cache.topn_key("process", 4)


def test_load_missing_sidecar_returns_none(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    assert setup_cache.load(data, [data]) is None


def test_load_corrupt_json_returns_none(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    setup_cache.sidecar_path(data).write_text("not valid json {{{")
    assert setup_cache.load(data, [data]) is None


def test_load_invalidates_on_dims_mismatch(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]
    setup_cache.save(data, files, SetupCache.empty(files))

    path = setup_cache.sidecar_path(data)
    raw = json.loads(path.read_text())
    raw["dims"]["K_MAX"] = raw["dims"]["K_MAX"] + 1
    path.write_text(json.dumps(raw))

    assert setup_cache.load(data, files) is None


def test_load_invalidates_on_format_version_mismatch(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]
    setup_cache.save(data, files, SetupCache.empty(files))

    path = setup_cache.sidecar_path(data)
    raw = json.loads(path.read_text())
    raw["format_version"] = raw["format_version"] + 1
    path.write_text(json.dumps(raw))

    assert setup_cache.load(data, files) is None


def test_load_invalidates_on_file_content_change(tmp_path):
    data = tmp_path / "shard.parquet"
    _touch_parquet(data, n=1)
    files = [data]
    setup_cache.save(data, files, SetupCache.empty(files))

    _touch_parquet(data, n=50)  # different size -> fingerprint changes

    assert setup_cache.load(data, files) is None


def test_load_returns_none_on_malformed_cache_body(tmp_path):
    """format_version/dims/fingerprint all check out, but the cache body
    itself doesn't match SetupCache.from_json's expected shape (e.g. hand-
    edited or written by a version that changed a nested key) — a clean
    miss, not a crash."""
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]
    setup_cache.save(data, files, SetupCache.empty(files))

    path = setup_cache.sidecar_path(data)
    raw = json.loads(path.read_text())
    raw["vocab"] = {"pdg_map": {"11": 0}}  # missing required "mat_map" key
    path.write_text(json.dumps(raw))

    echoed = []
    assert setup_cache.load(data, files, echo=echoed.append) is None
    assert any("malformed" in m for m in echoed)


def test_load_soft_warns_on_git_hash_mismatch_but_still_hits(tmp_path, capsys):
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]
    cache = SetupCache.empty(files)
    cache.git_hash = "not-a-real-git-hash"
    setup_cache.save(data, files, cache)

    loaded = setup_cache.load(data, files)

    assert loaded is not None
    err = capsys.readouterr().err
    assert "not-a-real-git-hash" in err


# ── save: atomicity / robustness ────────────────────────────────────────


def test_save_is_atomic_no_stray_tmp_file(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]
    setup_cache.save(data, files, SetupCache.empty(files))

    leftovers = [p for p in tmp_path.iterdir() if ".tmp." in p.name]
    assert leftovers == []


def test_save_degrades_gracefully_on_permission_error(tmp_path):
    if os.geteuid() == 0:
        pytest.skip("root bypasses directory permission bits")
    data_dir = tmp_path / "ro"
    data_dir.mkdir()
    data = _touch_parquet(data_dir / "shard.parquet")
    files = [data]

    warnings = []
    mode = data_dir.stat().st_mode
    data_dir.chmod(stat.S_IREAD | stat.S_IEXEC)
    try:
        setup_cache.save(data, files, SetupCache.empty(files), echo=warnings.append)
    finally:
        data_dir.chmod(mode)

    assert any("could not write" in w for w in warnings)
    assert not setup_cache.sidecar_path(data).exists()


def test_save_merges_non_colliding_normalizer_keys(tmp_path):
    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]

    cache1 = SetupCache.empty(files)
    cache1.normalizers["k1"] = _entry(n_train_steps=1)
    setup_cache.save(data, files, cache1)

    cache2 = SetupCache.empty(files)
    cache2.normalizers["k2"] = _entry(n_train_steps=2)
    setup_cache.save(data, files, cache2)

    loaded = setup_cache.load(data, files)
    assert loaded is not None
    assert set(loaded.normalizers.keys()) == {"k1", "k2"}
    assert loaded.normalizers["k1"].n_train_steps == 1
    assert loaded.normalizers["k2"].n_train_steps == 2


def test_save_is_serialized_against_concurrent_writers(tmp_path):
    """Without the flock in setup_cache.save(), two concurrent writers can
    both load() the same base state and merge their own section in
    independently, so whichever os.replace() lands last silently drops the
    other's key — a lost-update race, not a corrupt file. Each of these
    threads writes a distinct normalizer key many times over; if the
    load-merge-write critical section isn't actually serialized, at least
    one thread's key is likely to go missing from the final merged cache."""
    import threading

    data = _touch_parquet(tmp_path / "shard.parquet")
    files = [data]
    setup_cache.save(data, files, SetupCache.empty(files))

    n_writers, n_rounds = 6, 15

    def _writer(idx: int) -> None:
        for r in range(n_rounds):
            cache = SetupCache.empty(files)
            cache.normalizers[f"k{idx}"] = _entry(n_train_steps=r)
            setup_cache.save(data, files, cache)

    threads = [threading.Thread(target=_writer, args=(i,)) for i in range(n_writers)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    loaded = setup_cache.load(data, files)
    assert loaded is not None
    assert set(loaded.normalizers.keys()) == {f"k{i}" for i in range(n_writers)}
    for i in range(n_writers):
        assert loaded.normalizers[f"k{i}"].n_train_steps == n_rounds - 1


# ── energy_quantiles_from_sample / energy_quantile_at ───────────────────


def test_energy_quantiles_from_sample_empty():
    result = setup_cache.energy_quantiles_from_sample(np.empty(0, dtype=np.float32))
    assert result.size == 0


def test_energy_quantiles_from_sample_has_fixed_grid_size():
    sample = np.random.default_rng(0).normal(size=5000).astype(np.float32)
    result = setup_cache.energy_quantiles_from_sample(sample)
    assert result.shape == (setup_cache.ENERGY_QUANTILE_LEVELS,)
    assert result[0] == pytest.approx(sample.min(), abs=1e-3)
    assert result[-1] == pytest.approx(sample.max(), abs=1e-3)


def test_energy_quantile_at_matches_direct_quantile_on_stored_grid():
    sample = np.random.default_rng(1).exponential(size=20_000).astype(np.float32)
    grid = setup_cache.energy_quantiles_from_sample(sample)

    levels = np.linspace(0.0, 1.0, 5)
    got = setup_cache.energy_quantile_at(grid, levels)
    expected = np.quantile(sample, levels)

    np.testing.assert_allclose(got, expected, rtol=0.05)


def test_energy_quantile_at_median_of_two_points():
    grid = np.array([0.0, 10.0], dtype=np.float32)
    result = setup_cache.energy_quantile_at(grid, np.array([0.0, 0.5, 1.0]))
    np.testing.assert_allclose(result, [0.0, 5.0, 10.0])


# ── n_train_steps_for_split ──────────────────────────────────────────────


def test_n_train_steps_for_split_matches_full_scan():
    unique_ids = np.array([1, 2, 3, 4, 5])
    counts = np.array([10, 20, 30, 40, 50])
    train_events_arr = np.array([2, 4, 5])

    result = setup_cache.n_train_steps_for_split(unique_ids, counts, train_events_arr)

    assert result == 20 + 40 + 50


# ── compute_event_index_from_files: cross-file event_id offsetting ─────────


def test_compute_event_index_from_files_offsets_colliding_ids(tmp_path):
    """Two files that each restart event_id from 0 (one Geant4 job per file)
    must not have their same-numbered events collapsed into one by
    np.unique — each file's ids get shifted by a distinct offset first (see
    giant.data.loader.event_id_offset)."""
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path_a)
    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path_b)

    unique_ids, counts = setup_cache.compute_event_index_from_files([path_a, path_b])

    assert len(unique_ids) == 6
    assert int(counts.sum()) == 6
    assert np.all(counts == 1)


def test_compute_event_index_from_files_single_file_unaffected(tmp_path):
    """A single file's ids are offset by 0 (event_id_offset(0) == 0), so a
    single-file load's unique ids/counts are unchanged by the offsetting."""
    path = tmp_path / "a.parquet"
    pd.DataFrame({"event_id": [5, 5, 7]}).to_parquet(path)

    unique_ids, counts = setup_cache.compute_event_index_from_files([path])

    np.testing.assert_array_equal(unique_ids, [5, 7])
    np.testing.assert_array_equal(counts, [2, 1])


def test_compute_event_index_from_files_empty_file_list():
    unique_ids, counts = setup_cache.compute_event_index_from_files([])
    assert unique_ids.size == 0
    assert counts.size == 0
    assert unique_ids.dtype == np.int64
    assert counts.dtype == np.int64