import pytest
import torch

import kernels

lp = kernels.get_kernel("phanerozoic/logits-processor", version=1, trust_remote_code=True)

requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")


def _pack(allowed_sets, V, device):
    nwords = (V + 31) // 32
    mask = torch.zeros(len(allowed_sets), nwords, dtype=torch.int32, device=device)
    for r, s in enumerate(allowed_sets):
        for tok in s:
            mask[r, tok >> 5] |= (1 << (tok & 31))
    return mask


@requires_cuda
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_bitmask_sets_disallowed_to_neg_inf(dtype):
    V = 130
    logits = torch.randn(2, V, device="cuda", dtype=dtype)
    allowed = [{3, 7, 50, 129}, {0, 42, 63}]
    ref = logits.clone()
    lp.apply_token_bitmask(logits, _pack(allowed, V, "cuda"))
    for r, s in enumerate(allowed):
        for tok in range(V):
            if tok in s:
                assert logits[r, tok] == ref[r, tok]
            else:
                assert logits[r, tok].float() == float("-inf")


@requires_cuda
def test_greedy_and_topk1_are_argmax():
    logits = torch.randn(8, 2000, device="cuda")
    amax = logits.argmax(-1)
    assert torch.equal(lp.sample(logits, temperature=0.0), amax)
    assert torch.equal(lp.sample(logits, top_k=1, temperature=1.0), amax)


@requires_cuda
def test_min_p_and_topp_exclude_low_prob():
    logits = torch.full((4, 50), -10.0, device="cuda")
    logits[:, 7] = 10.0  # one dominant token, softmax ~1.0
    assert torch.equal(lp.sample(logits, temperature=1.0, min_p=0.5, seed=1),
                       torch.full((4,), 7, device="cuda"))
    assert torch.equal(lp.sample(logits, temperature=1.0, top_p=0.9, seed=2),
                       torch.full((4,), 7, device="cuda"))


@requires_cuda
def test_deterministic():
    logits = torch.randn(16, 500, device="cuda")
    a = lp.sample(logits, temperature=1.0, seed=99, offset=3)
    b = lp.sample(logits, temperature=1.0, seed=99, offset=3)
    assert torch.equal(a, b)


@requires_cuda
def test_distribution_matches_softmax():
    torch.manual_seed(0)
    V, N = 32, 40000
    base = torch.randn(V, device="cuda")
    logits = base.unsqueeze(0).repeat(N, 1)
    toks = lp.sample(logits, temperature=1.0, seed=1234)
    emp = torch.bincount(toks, minlength=V).float() / N
    ref = torch.softmax(base, -1)
    assert (emp - ref).abs().max().item() < 0.03