import pytest
import torch

import kernels

bitnet = kernels.get_kernel("phanerozoic/bitnet-cpu", version=1, trust_remote_code=True)


def unpack_ternary(wp):
    cols = [((wp >> (2 * j)) & 3).to(torch.int16) - 2 for j in range(4)]
    return torch.stack(cols, dim=-1).reshape(wp.shape[0], wp.shape[1] * 4)


@pytest.mark.kernels_ci
@pytest.mark.parametrize("M", [1, 4, 16, 128])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
def test_gemm_matches_exact_reference(M, dtype):
    """The integer path is exact: feeding the kernel's own quantized operands
    into an f32 reference bounds the residual by bf16 output rounding alone
    (one ulp, 2^-8 relative). An independent Python re-quantization is not a
    valid reference: dividing by amax/127 and multiplying by 127/amax differ
    by one ulp at rounding boundaries and flip occasional codes."""
    torch.manual_seed(0)
    N, K = 512, 2560
    W = torch.randint(-1, 2, (N, K), dtype=torch.int8)
    wp = bitnet.pack_weights(W)
    sw = (torch.rand(N) * 0.5 + 0.5).to(torch.bfloat16)
    x = torch.randn(M, K, dtype=dtype)

    q, s = bitnet.quantize_activation(x)
    y = bitnet.bitnet_gemm(q, wp, s, sw).float()
    ref = (q.float() @ unpack_ternary(wp).float().t()) * s.float().unsqueeze(-1) * sw.float().unsqueeze(0)
    rel = ((y - ref).abs() / ref.abs().clamp(min=1.0)).max().item()
    assert rel < 8e-3, f"max rel {rel}"


@pytest.mark.kernels_ci
@pytest.mark.parametrize("M", [1, 4, 15])
def test_fused_path_matches_gemm_path(M):
    """The fused (M<16) path quantizes internally with the same code as
    quantize_activation; outputs agree to bf16 rounding of the scale."""
    torch.manual_seed(1)
    N, K = 1024, 4096
    W = torch.randint(-1, 2, (N, K), dtype=torch.int8)
    wp = bitnet.pack_weights(W)
    sw = torch.ones(N, dtype=torch.bfloat16)
    x = torch.randn(M, K, dtype=torch.bfloat16)
    y_fused = bitnet.bitnet_gemv_fused(x, wp, sw).float()
    q, s = bitnet.quantize_activation(x)
    y_split = bitnet.bitnet_gemm(q, wp, s, sw).float()
    rel = ((y_fused - y_split).abs() / y_split.abs().clamp(min=1.0)).max().item()
    assert rel < 8e-3, f"max rel {rel}"


@pytest.mark.kernels_ci
def test_quantize_activation_roundtrip():
    torch.manual_seed(2)
    x = torch.randn(8, 1024, dtype=torch.bfloat16)
    q, s = bitnet.quantize_activation(x)
    assert q.dtype == torch.int8 and s.dtype == torch.bfloat16
    assert (q.abs() <= 127).all()
    recon = q.float() * s.float().unsqueeze(-1)
    torch.testing.assert_close(recon, x.float(), rtol=2e-2, atol=2e-2)


@pytest.mark.kernels_ci
def test_bitlinear_module():
    torch.manual_seed(3)
    lin = torch.nn.Linear(2560, 512, bias=False)
    bl = bitnet.BitLinear.from_dense(lin)
    x = torch.randn(4, 2560, dtype=torch.bfloat16)
    y = bl(x)
    assert y.shape == (4, 512) and y.dtype == torch.bfloat16
    assert torch.isfinite(y.float()).all()