import torch
from giant.constants import COND_DIM
from giant.model.network import DenoisingMLP, SinusoidalEmbedding


def test_sinusoidal_embedding_shape():
    emb = SinusoidalEmbedding(64)
    t = torch.rand(16)
    assert emb(t).shape == (16, 64)


def test_sinusoidal_embedding_batch_1():
    emb = SinusoidalEmbedding(32)
    t = torch.tensor([0.5])
    assert emb(t).shape == (1, 32)


def test_denoising_mlp_output_shape():
    B = 8
    model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
    x_t = torch.randn(B, 9)
    t = torch.rand(B)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack(
        [
            torch.randint(0, 5, (B,)),
            torch.randint(0, 3, (B,)),
        ],
        dim=1,
    )
    out = model(x_t, t, cond_cont, cond_cat)
    assert out.shape == (B, 9)


def test_denoising_mlp_gradients_flow():
    B = 4
    model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
    x_t = torch.randn(B, 9)
    t = torch.rand(B)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    loss = model(x_t, t, cond_cont, cond_cat).sum()
    loss.backward()
    for name, p in model.named_parameters():
        assert p.grad is not None, f"no grad for {name}"