import numpy as np
import torch

from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.validate import validate_marginals


def _tiny_models():
    s1 = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=16, n_blocks=1)
    s2 = SecondaryDecoder(pdg_vocab=3, mat_vocab=2, hidden_dim=16, n_blocks=1)
    return s1.eval(), s2.eval()


def _zero_secondaries_loader(B=4, n_batches=2):
    """A val_loader whose every batch has n_sec=0 (real side) — matches the
    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) tuple shape
    StreamingStepsDataset yields."""
    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.zeros(B, dtype=torch.long)
        sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM)
        proc_idx = torch.zeros(B, dtype=torch.long)
        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_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 -- a regression the old species/bincount code this
    replaced explicitly guarded against."""
    s1, s2 = _tiny_models()
    loader = _zero_secondaries_loader()

    # Force the Stage-1 n_sec head's prediction to 0 for every sample too, so
    # the generated side's valid-slot mask is also empty (real side is
    # already all n_sec=0 by construction of the fake loader above).
    def _fake_sample_flow(model, cond_cont, cond_cat, **kw):
        B = cond_cont.size(0)
        return torch.randn(B, X_DIM), torch.zeros(B, dtype=torch.long)

    monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow)

    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=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()