import kernels
import pytest
import torch

k = kernels.get_kernel("GoedelMachines/dg-w4-kernels")


@pytest.mark.kernels_ci
def test_fused_entropy_matches_categorical():
    x = torch.randn(64, 8192, device="cuda", dtype=torch.float32) * 3
    ref = torch.distributions.Categorical(logits=x).entropy()
    got = k.fused_entropy(x)
    torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4)


@pytest.mark.kernels_ci
def test_gumbel_argmax_returns_exact_argmax():
    x = torch.randn(32, 4096, device="cuda", dtype=torch.float32)
    _, amax = k.gumbel_argmax_sample(x, seed=1234)
    assert (amax == x.argmax(-1)).all()


@pytest.mark.kernels_ci
def test_w4a16_matches_dequant_reference():
    W = torch.randn(512, 1024, device="cuda") * 0.02
    qw, s, z = k.quantize_w4(W, 128)
    Wdq = k.dequant_w4(qw, s, z, 128)
    x = torch.randn(64, 1024, device="cuda", dtype=torch.float16)
    ref = torch.nn.functional.linear(x, Wdq)
    got = k.w4a16_linear(x, qw, s, z, BK=128)
    rel = (got.float() - ref.float()).abs().max() / ref.float().abs().max()
    assert rel < 2e-2, rel


@pytest.mark.kernels_ci
def test_fused_rmsnorm_matches_reference():
    x = torch.randn(128, 2816, device="cuda", dtype=torch.bfloat16)
    w = torch.randn(2816, device="cuda", dtype=torch.bfloat16)
    ref = (x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-6) * w.float())
    got = k.fused_rmsnorm(x, w, 1e-6).float()
    assert ((got - ref).abs().mean() / ref.abs().mean()) < 5e-3