"""Both ops against `gguf.quants.dequantize`, the reference implementation of the block layouts.

The blocks are random bytes rather than a quantized tensor: `gguf` can unpack every type but only
pack a couple of them, and unpacking is defined for any byte pattern, so this needs no quantizer and
no checkpoint. The one constraint is that a block's scales are fp16 fields — masked below so a random
pattern cannot land on an exponent of all ones and make the whole block inf/nan.
"""

import numpy as np
import pytest
import torch

from gguf_kernels import MAX_GEMV_ROWS, dequantize, mul_mat_vec


gguf = pytest.importorskip("gguf", reason="the reference unpacker comes from the `gguf` package")

# name -> (ggml type id, values per block, bytes per block)
QUANT_TYPES = {
    "Q4_K": (12, 256, 144),
    "Q5_K": (13, 256, 176),
    "Q6_K": (14, 256, 210),
    "Q8_0": (8, 32, 34),
}


def random_blocks(rows: int, cols: int, ggml_name: str, device="cuda"):
    """Random blocks and the values `gguf` reads out of them."""
    _, block_values, block_bytes = QUANT_TYPES[ggml_name]
    generator = np.random.default_rng(0)
    packed = generator.integers(0, 256, (rows, cols // block_values * block_bytes), dtype=np.uint8)
    # every fp16 scale sits at an even offset in its block, so clearing bit 6 of each odd byte keeps
    # every possible fp16 field finite whatever the rest of the pattern is
    packed[:, 1::2] &= 0xBF

    quant_type = getattr(gguf.GGMLQuantizationType, ggml_name)
    reference = gguf.quants.dequantize(packed.reshape(-1), quant_type).reshape(rows, cols)
    return torch.from_numpy(packed).to(device), torch.from_numpy(reference).to(device)


@pytest.mark.kernels_ci
@pytest.mark.parametrize("ggml_name", QUANT_TYPES)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_dequantize_matches_reference(ggml_name, dtype):
    ggml_type = QUANT_TYPES[ggml_name][0]
    rows, cols = 64, 512
    blocks, reference = random_blocks(rows, cols, ggml_name)

    out = dequantize(blocks, ggml_type, rows, cols, dtype)

    assert out.shape == (rows, cols) and out.dtype == dtype
    # the kernel writes `dtype` directly, so the tolerance is that dtype's own resolution
    torch.testing.assert_close(out.float(), reference, rtol=torch.finfo(dtype).eps * 4, atol=0)


@pytest.mark.kernels_ci
@pytest.mark.parametrize("ggml_name", QUANT_TYPES)
@pytest.mark.parametrize("n_rows", [1, MAX_GEMV_ROWS])
def test_mul_mat_vec_matches_matmul(ggml_name, n_rows):
    ggml_type = QUANT_TYPES[ggml_name][0]
    out_features, in_features = 128, 512
    blocks, reference = random_blocks(out_features, in_features, ggml_name)
    x = torch.randn(n_rows, in_features, dtype=torch.bfloat16, device="cuda")

    out = mul_mat_vec(blocks, x, ggml_type, out_features)

    assert out.shape == (n_rows, out_features) and out.dtype == torch.float32
    # the kernel quantizes the activations to q8_1, so this is close to a matmul, not equal to one
    expected = x.float() @ reference.T
    torch.testing.assert_close(out, expected, rtol=2e-2, atol=2e-2 * expected.abs().max())


@pytest.mark.kernels_ci
def test_gemv_is_compileable():
    """A graph break here would cost more than the kernel saves, so the fake has to be right."""
    ggml_type = QUANT_TYPES["Q4_K"][0]
    out_features, in_features = 128, 512
    blocks, reference = random_blocks(out_features, in_features, "Q4_K")
    x = torch.randn(1, in_features, dtype=torch.bfloat16, device="cuda")

    compiled = torch.compile(
        lambda t: mul_mat_vec(blocks, t, ggml_type, out_features), fullgraph=True
    )

    expected = x.float() @ reference.T
    torch.testing.assert_close(compiled(x), expected, rtol=2e-2, atol=2e-2 * expected.abs().max())