import torch
from giant.constants import COND_DIM
from giant.model.network import Stage1Model
from giant.model.schedule import CosineSchedule, flow_matching_loss
from giant.sample import sample_flow, sample_ddim

PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}


def _small_model():
    return Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=PARTICLE_CFG,
        material_cfg=MATERIAL_CFG,
        hidden_dim=32,
        n_res_blocks=2,
        n_sec_head_k_max=15,
    )


def _batch(B=8):
    x1 = torch.randn(B, 9)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    return x1, cond_cont, cond_cat


def test_flow_matching_loss_nonneg():
    x1, cond_cont, cond_cat = _batch()
    loss = flow_matching_loss(_small_model(), x1, cond_cont, cond_cat)
    assert loss.item() >= 0.0


def test_flow_matching_loss_is_scalar():
    x1, cond_cont, cond_cat = _batch()
    loss = flow_matching_loss(_small_model(), x1, cond_cont, cond_cat)
    assert loss.shape == ()


def test_flow_matching_loss_has_grad():
    model = _small_model()
    x1, cond_cont, cond_cat = _batch()
    flow_matching_loss(model, x1, cond_cont, cond_cat).backward()
    assert any(p.grad is not None for p in model.parameters())


def test_sample_flow_shape():
    B = 6
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
    assert sample.shape == (B, 9)
    assert n_sec is not None and n_sec.shape == (B,)


def test_ddpm_loss_nonneg():
    schedule = CosineSchedule(T=50)
    x1, cond_cont, cond_cat = _batch()
    loss = schedule.loss(_small_model(), x1, cond_cont, cond_cat)
    assert loss.item() >= 0.0


def test_sample_ddim_shape():
    B = 4
    schedule = CosineSchedule(T=50)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
    assert sample.shape == (B, 9)
    assert n_sec is not None and n_sec.shape == (B,)