import pytest
import torch
from torch.nn.utils.rnn import pad_sequence
from kernels.varlen_packing import pack_varlen_batch, unpack_varlen_batch, packed_to_padded


def test_roundtrip(device):
    torch.manual_seed(42)
    D    = 768
    lens = [160, 80, 90, 110, 140, 70, 130, 100]
    toks = [torch.randn(L, D, device=device) for L in lens]
    packed, cu = pack_varlen_batch(toks)
    recovered  = unpack_varlen_batch(packed, cu)
    for i, (orig, rec) in enumerate(zip(toks, recovered)):
        assert torch.allclose(orig, rec), f"Mismatch at item {i}"


def test_cu_seqlens(device):
    D    = 64
    lens = [10, 20, 15]
    toks = [torch.randn(L, D, device=device) for L in lens]
    _, cu = pack_varlen_batch(toks)
    expected = torch.tensor([0, 10, 30, 45], dtype=torch.int32, device=device)
    assert torch.equal(cu, expected)


def test_less_memory_than_padded(device):
    D    = 768
    lens = [196, 40, 50, 60]
    toks = [torch.randn(L, D, device=device) for L in lens]
    packed, _ = pack_varlen_batch(toks)
    padded    = pad_sequence(toks, batch_first=True)
    assert packed.numel() < padded.numel()


def test_attention_mask(device):
    D    = 64
    lens = [5, 3, 4]
    toks = [torch.randn(L, D, device=device) for L in lens]
    packed, cu = pack_varlen_batch(toks)
    padded, mask = packed_to_padded(packed, cu)
    assert padded.shape == (3, 5, D)
    assert mask[0].all()
    assert mask[1, :3].all()
    assert not mask[1, 3:].any()


def test_single_item(device):
    D = 256
    toks = [torch.randn(100, D, device=device)]
    packed, cu = pack_varlen_batch(toks)
    assert packed.shape == (100, D)
    assert cu.tolist() == [0, 100]