giant/tests/test_checkpoint_io.py

"""Tests for giant.checkpoint_io.load_for_inference (issues.md Issue 5) —
the shared bootstrap `giant predict`/`giant rollout` use to go from a
checkpoint path to ready-to-run models."""

from __future__ import annotations

import copy

import numpy as np
import pytest
import torch

from giant import config as gconfig
from giant.checkpoint_io import (
    CheckpointCompatibilityError,
    InferenceContext,
    apply_config_overrides,
    conditioning_axes,
    load_for_inference,
    stage_cfg,
)
from giant.data.loader import TopNMap
from giant.data.setup_cache import topnmap_to_json
from giant.data.transforms import Normalizer
from giant.model.network import build_models

PDG_MAP = {11: 0, 22: 1, -11: 2}
MAT_MAP = {"G4_PbWO4": 0, "G4_AIR": 1}


def _model_cfg(stage2_active: bool = True) -> dict:
    """DEFAULT_CONFIG-derived, shrunk for speed — same pattern as
    tests/test_network.py::_minimal_model_config. Default `conditioning`
    (both axes "physical") needs no top-N vocab map, so this is a cheap,
    fully self-contained happy-path config."""
    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
    cfg["conditioning"]["particle"]["emb_dim"] = 4
    cfg["conditioning"]["material"]["emb_dim"] = 4
    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
    cfg["stage2_model"]["active"] = stage2_active
    return {
        "pdg_vocab": len(PDG_MAP),
        "mat_vocab": len(MAT_MAP),
        "conditioning": cfg["conditioning"],
        "stage1_model": cfg["stage1_model"],
        "stage2_model": cfg["stage2_model"],
    }


def _norms() -> tuple[Normalizer, Normalizer, Normalizer]:
    rng = np.random.default_rng(0)
    cond = Normalizer().fit(rng.standard_normal((100, 15)).astype(np.float32))
    tgt = Normalizer().fit(rng.standard_normal((100, 9)).astype(np.float32))
    sec_phys = Normalizer().fit(rng.standard_normal((100, 2)).astype(np.float32))
    return cond, tgt, sec_phys


def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overrides):
    cfg = model_cfg if model_cfg is not None else _model_cfg()
    built = build_models(cfg)
    stage1, stage2 = built["stage1"], built["stage2"]
    cond, tgt, sec_phys = _norms()
    ckpt: dict = {
        "model_config": cfg,
        "model": stage1.state_dict() if stage1 is not None else {},
        "sec_decoder": stage2.state_dict() if stage2 is not None else {},
        "pdg_map": PDG_MAP,
        "mat_map": MAT_MAP,
        "normalizer": {"cond": cond.to_dict(), "target": tgt.to_dict(), "sec_phys": sec_phys.to_dict()},
        "epoch": 3,
        "best_val_loss": 0.5,
    }
    if ema:
        ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
        ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
    # DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
    # "onehot", and giant train's pipeline (gitea #29) now always writes a
    # sec_type_topn_map in that case — default one in here too, unless a
    # test explicitly overrides it, so fixtures represent a real, loadable
    # checkpoint by default rather than exercising the "missing" guard by
    # accident.
    particle_type_target = cfg.get("stage2_model", {}).get("particle_type", {}).get("target", "onehot")
    if particle_type_target == "onehot" and "sec_type_topn_map" not in ckpt_overrides:
        default_sec_type_topn = TopNMap(class_map=dict(zip(PDG_MAP, range(len(PDG_MAP)))), other_members={})
        ckpt["sec_type_topn_map"] = topnmap_to_json(default_sec_type_topn)
    ckpt.update(ckpt_overrides)
    path = tmp_path / "ckpt.pt"
    torch.save(ckpt, path)
    return path


def _onehot_model_cfg() -> dict:
    cfg = _model_cfg()
    cfg["conditioning"]["particle"]["type"] = "onehot"
    return cfg


# ---------------------------------------------------------------------------
# Happy path
# ---------------------------------------------------------------------------


def test_happy_path_returns_populated_context(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")

    assert isinstance(ctx, InferenceContext)
    assert ctx.stage1 is not None and ctx.stage2 is not None
    assert not ctx.stage1.training
    assert not ctx.stage2.training
    assert next(ctx.stage1.parameters()).device == torch.device("cpu")
    assert ctx.pdg_map == PDG_MAP
    assert ctx.mat_map == MAT_MAP
    assert all(isinstance(k, int) for k in ctx.pdg_map)
    assert all(isinstance(k, str) for k in ctx.mat_map)
    assert ctx.particle_conditioning == "physical"
    assert ctx.material_conditioning == "physical"
    assert ctx.k_max == 3
    assert ctx.epoch == 3
    assert ctx.best_val_loss == 0.5
    assert ctx.model_config["stage1_model"]["hidden_dim"] == 8


def test_happy_path_normalizer_values_round_trip(tmp_path):
    cond, tgt, sec_phys = _norms()
    checkpoint = _write_checkpoint(tmp_path)
    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")

    assert ctx.cond_norm.mean is not None and cond.mean is not None
    assert ctx.tgt_norm.mean is not None and tgt.mean is not None
    assert ctx.sec_phys_norm.mean is not None and sec_phys.mean is not None
    np.testing.assert_allclose(ctx.cond_norm.mean, cond.mean)
    np.testing.assert_allclose(ctx.tgt_norm.mean, tgt.mean)
    np.testing.assert_allclose(ctx.sec_phys_norm.mean, sec_phys.mean)


# ---------------------------------------------------------------------------
# Guards
# ---------------------------------------------------------------------------


def test_missing_model_config_raises(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ckpt = torch.load(checkpoint, weights_only=False)
    del ckpt["model_config"]
    torch.save(ckpt, checkpoint)

    with pytest.raises(CheckpointCompatibilityError, match="no model_config"):
        load_for_inference(checkpoint, torch.device("cpu"), "predict")


def test_missing_sec_decoder_raises(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ckpt = torch.load(checkpoint, weights_only=False)
    del ckpt["sec_decoder"]
    torch.save(ckpt, checkpoint)

    with pytest.raises(CheckpointCompatibilityError, match="no sec_decoder"):
        load_for_inference(checkpoint, torch.device("cpu"), "predict")


def test_missing_sec_phys_normalizer_raises(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ckpt = torch.load(checkpoint, weights_only=False)
    del ckpt["normalizer"]["sec_phys"]
    torch.save(ckpt, checkpoint)

    with pytest.raises(CheckpointCompatibilityError, match="no normalizer.sec_phys"):
        load_for_inference(checkpoint, torch.device("cpu"), "predict")


def test_onehot_particle_conditioning_without_topn_map_raises(tmp_path):
    checkpoint = _write_checkpoint(tmp_path, model_cfg=_onehot_model_cfg())

    with pytest.raises(CheckpointCompatibilityError, match="pdg_topn_map"):
        load_for_inference(checkpoint, torch.device("cpu"), "predict")


def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
    topn = TopNMap(class_map={11: 0, 22: 1}, other_members={})
    checkpoint = _write_checkpoint(
        tmp_path,
        model_cfg=_onehot_model_cfg(),
        pdg_topn_map=topnmap_to_json(topn),
    )
    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
    assert ctx.particle_conditioning == "onehot"
    assert ctx.pdg_topn_map is not None
    assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}


def test_onehot_particle_type_target_without_sec_type_topn_map_raises(tmp_path):
    """DEFAULT_CONFIG's stage2_model.particle_type.target="onehot" needs a
    sec_type_topn_map (gitea #29) — a checkpoint with neither key at all
    (not even the pre-#29 pdg_topn_map to fall back to) must fail loudly."""
    checkpoint = _write_checkpoint(tmp_path, sec_type_topn_map=None)
    ckpt = torch.load(checkpoint, weights_only=False)
    del ckpt["sec_type_topn_map"]
    torch.save(ckpt, checkpoint)

    with pytest.raises(CheckpointCompatibilityError, match="sec_type_topn_map"):
        load_for_inference(checkpoint, torch.device("cpu"), "predict")


def test_pre_gitea_29_checkpoint_falls_back_to_pdg_topn_map_for_sec_type(tmp_path):
    """A checkpoint written before gitea #29 has no sec_type_topn_map key at
    all — conditioning and secondary-type onehot maps were always the same
    map, saved once under pdg_topn_map. load_for_inference must reproduce
    that exact pre-#29 behavior for such a checkpoint."""
    topn = TopNMap(class_map={11: 0, 22: 1, -11: 2}, other_members={})
    checkpoint = _write_checkpoint(
        tmp_path,
        model_cfg=_onehot_model_cfg(),
        pdg_topn_map=topnmap_to_json(topn),
        sec_type_topn_map=None,
    )
    ckpt = torch.load(checkpoint, weights_only=False)
    del ckpt["sec_type_topn_map"]
    torch.save(ckpt, checkpoint)

    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
    assert ctx.sec_type_topn_map is not None
    assert ctx.sec_type_topn_map.class_map == {11: 0, 22: 1, -11: 2}


def test_ema_weights_requested_but_missing_raises(tmp_path):
    checkpoint = _write_checkpoint(tmp_path, ema=False)

    with pytest.raises(CheckpointCompatibilityError, match="no EMA weights"):
        load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")


def test_ema_weights_requested_and_present_succeeds(tmp_path):
    checkpoint = _write_checkpoint(tmp_path, ema=True)
    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
    assert ctx.stage1 is not None and ctx.stage2 is not None


@pytest.mark.parametrize("command_name", ["predict", "rollout"])
def test_inactive_stage_with_require_stage2_raises_with_command_name(tmp_path, command_name):
    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))

    with pytest.raises(CheckpointCompatibilityError, match=f"{command_name} needs both"):
        load_for_inference(checkpoint, torch.device("cpu"), command_name)


def test_inactive_stage_with_require_stage2_false_succeeds_with_stage2_none(tmp_path):
    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))

    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", require_stage2=False)
    assert ctx.stage1 is not None
    assert ctx.stage2 is None


# ---------------------------------------------------------------------------
# conditioning_axes / stage_cfg
# ---------------------------------------------------------------------------


def test_conditioning_axes_v02_flat_string_applies_to_both_axes():
    assert conditioning_axes({"conditioning": "embedding"}) == ("embedding", "embedding")


def test_conditioning_axes_v03_nested_dict_independent_per_axis():
    model_cfg = {"conditioning": {"particle": {"type": "onehot"}, "material": {"type": "physical"}}}
    assert conditioning_axes(model_cfg) == ("onehot", "physical")


def test_conditioning_axes_missing_key_uses_default():
    assert conditioning_axes({}, default="embedding") == ("embedding", "embedding")


def test_stage_cfg_new_shape_returns_subdict():
    model_cfg = {"stage2_model": {"k_max": 7}}
    assert stage_cfg(model_cfg, "stage2") == {"k_max": 7}


def test_stage_cfg_v02_flat_shape_returns_empty_dict():
    model_cfg = {"hidden_dim": 32, "n_blocks": 4}
    assert stage_cfg(model_cfg, "stage2") == {}


# ---------------------------------------------------------------------------
# config_overrides (gitea #87)
# ---------------------------------------------------------------------------


def _router_model_cfg() -> dict:
    cfg = _model_cfg()
    cfg["stage1_model"]["router"] = {"enabled": True, "type": "energy", "n_experts": 2}
    return cfg


def test_config_override_n_sec_sampling_changes_stage2_attribute(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ctx = load_for_inference(
        checkpoint,
        torch.device("cpu"),
        "predict",
        config_overrides={"stage2_model.n_sec.sampling": "sample"},
    )
    assert ctx.stage2 is not None
    assert ctx.stage2.n_sec_sampling == "sample"
    assert ctx.config_overrides == {"stage2_model.n_sec.sampling": "sample"}


def test_config_override_ddpm_n_steps_changes_context_fields(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ctx = load_for_inference(
        checkpoint,
        torch.device("cpu"),
        "predict",
        config_overrides={"stage1_model.ddpm.n_steps": 42, "stage2_model.ddpm.n_steps": 7},
    )
    assert ctx.stage1_ddpm_steps == 42
    assert ctx.stage2_ddpm_steps == 7


def test_config_override_other_policy_changes_context_field(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ctx = load_for_inference(
        checkpoint,
        torch.device("cpu"),
        "predict",
        config_overrides={"stage2_model.particle_type.other_policy": "modal"},
    )
    assert ctx.other_policy == "modal"


def test_config_override_router_temperature_changes_router_attribute(tmp_path):
    checkpoint = _write_checkpoint(tmp_path, model_cfg=_router_model_cfg())
    ctx = load_for_inference(
        checkpoint,
        torch.device("cpu"),
        "predict",
        config_overrides={"stage1_model.router.temperature": 1.5},
    )
    assert ctx.stage1 is not None
    assert ctx.stage1.trunk.router.temperature == pytest.approx(1.5)


def test_config_override_no_overrides_defaults_to_empty_dict(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
    assert ctx.config_overrides == {}


def test_config_override_unknown_path_raises(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
        load_for_inference(
            checkpoint,
            torch.device("cpu"),
            "predict",
            config_overrides={"stage2_model.n_sec.typo": "sample"},
        )


def test_config_override_shape_bearing_key_raises_up_front(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
        load_for_inference(
            checkpoint,
            torch.device("cpu"),
            "predict",
            config_overrides={"stage1_model.hidden_dim": 999},
        )


def test_config_override_bad_value_raises(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    with pytest.raises(CheckpointCompatibilityError, match="must be one of"):
        load_for_inference(
            checkpoint,
            torch.device("cpu"),
            "predict",
            config_overrides={"stage2_model.n_sec.sampling": "maybe"},
        )


def test_apply_config_overrides_no_overrides_returns_same_object():
    cfg = _model_cfg()
    assert apply_config_overrides(cfg, None) is cfg
    assert apply_config_overrides(cfg, {}) is cfg


def test_apply_config_overrides_migrates_legacy_flat_model_config_first():
    legacy_cfg = {
        "pdg_vocab": len(PDG_MAP),
        "mat_vocab": len(MAT_MAP),
        "hidden_dim": 32,
        "n_blocks": 4,
        "emb_dim": 8,
        "dropout": 0.1,
        "k_max": 5,
    }
    merged = apply_config_overrides(legacy_cfg, {"stage1_model.ddpm.n_steps": 10})
    assert merged["stage1_model"]["ddpm"]["n_steps"] == 10
    assert merged["stage1_model"]["hidden_dim"] == 32