import numpy as np
import torch

from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim
from giant.validate import validate_marginals

_PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
_MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
_K_MAX = 5


def _tiny_models(particle_type_cfg: dict | None = None):
    """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head, so n_sec always
    comes from Stage2OneShot."""
    s1 = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PARTICLE_CFG,
        material_cfg=_MATERIAL_CFG,
        hidden_dim=16,
        n_res_blocks=1,
    )
    target = (particle_type_cfg or {}).get("target", "physical")
    sec_dim = stage2_trunk_sec_dim(
        particle_type_cfg or {"target": "physical"},
        "flow",
        _K_MAX,
        int(_PARTICLE_CFG["emb_dim"]),
    )
    s2 = Stage2OneShot(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PARTICLE_CFG,
        material_cfg=_MATERIAL_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        generator="flow",
        time_dim=16,
        k_max=_K_MAX,
        sec_dim=sec_dim,
        particle_type_cfg=particle_type_cfg,
    )
    assert s2.particle_type_cfg.get("target", "physical") == target
    return s1.eval(), s2.eval()


def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
    """A val_loader matching StreamingStepsDataset's 7-tuple batch shape:
    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)."""
    batches = []
    for _ in range(n_batches):
        cond_cont = torch.randn(B, COND_DIM)
        cond_cat = torch.zeros(B, 2, dtype=torch.long)
        x1 = torch.randn(B, X_DIM)
        n_sec = torch.full((B,), n_sec_value, dtype=torch.long)
        sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
        proc_idx = torch.zeros(B, dtype=torch.long)
        sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
        batches.append(
            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
        )
    return batches


def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch):
    """If n_sec_pred collapses to 0 across the whole validated set (realistic
    during early/unstable training), phys_kl must degrade to NaN instead of
    crashing on the empty-array .min()/.max() reduction inside
    _histogram_kl."""
    s1, s2 = _tiny_models()
    loader = _loader(n_sec_value=0)

    def _fake_resolve_n_sec(
        stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
    ):
        return torch.zeros(cond_cont.size(0), dtype=torch.long)

    monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)

    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2, steps=2)

    assert np.asarray(result["phys_real"]).shape == (0, 2)
    assert np.asarray(result["phys_generated"]).shape == (0, 2)
    assert np.isnan(np.asarray(result["phys_kl"])).all()


def test_validate_marginals_physical_target_shapes():
    s1, s2 = _tiny_models()
    loader = _loader(n_sec_value=2)

    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)

    assert np.asarray(result["real"]).shape == (4, X_DIM)
    assert np.asarray(result["generated"]).shape == (4, X_DIM)
    assert np.asarray(result["kl_divergence"]).shape == (X_DIM,)
    assert "phys_real" in result and "phys_generated" in result and "phys_kl" in result
    assert "type_class_real" not in result


def test_validate_marginals_onehot_type_class_marginal():
    particle_type_cfg = {"target": "onehot"}
    s1, s2 = _tiny_models(particle_type_cfg)
    loader = _loader(n_sec_value=2, n_classes=s2.type_dim)

    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)

    assert "phys_real" not in result
    # Real/generated valid-slot counts need not agree (real: ground-truth
    # n_sec=2 always; generated: the untrained n_sec_head's own prediction).
    assert np.asarray(result["type_class_real"]).ndim == 1
    assert np.asarray(result["type_class_gen"]).ndim == 1
    assert np.asarray(result["type_class_real"]).shape[0] > 0


def test_validate_marginals_without_sec_decoder_returns_stage1_only():
    s1, _ = _tiny_models()
    loader = _loader(n_sec_value=0)

    result = validate_marginals(s1, loader, n_batches=1, steps=2)

    assert set(result) == {"real", "generated", "kl_divergence"}