"""Self-consistency tests for ref_kernels_torch (the pure-torch stand-in for the reference kernel.py).

    TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I examples/neuron/test_ref_kernels_torch.py
"""

import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ref_kernels_torch as K  # noqa: E402


def rel(a, b):
    return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-30)).item()


def test_act_quant():
    x = torch.randn(7, 3, 256).bfloat16() * 3
    q, s = K.act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu)
    assert q.dtype == torch.float8_e4m3fn and s.dtype == torch.float8_e8m0fnu and s.shape == (7, 3, 8)
    sf = s.float()
    assert torch.all(torch.log2(sf) == torch.round(torch.log2(sf)))
    deq = (q.float().unflatten(-1, (8, 32)) * sf.unsqueeze(-1)).flatten(-2)
    y = x.clone()
    K.act_quant(y, 32, "ue8m0", torch.float8_e8m0fnu, True)
    assert torch.equal(y, deq.bfloat16())
    z = y.clone()
    K.act_quant(z, 32, "ue8m0", torch.float8_e8m0fnu, True)
    assert torch.equal(z, y), "quant->dequant must be idempotent"
    assert rel(y, x) < 0.05
    # every block amax maps below 448 after scaling, and the scale is the smallest power of two that allows it
    amax = x.float().unflatten(-1, (8, 32)).abs().amax(-1)
    assert torch.all(amax / sf <= 448) and torch.all(amax / (sf / 2) > 448)


def test_fp8_gemm():
    torch.set_default_dtype(torch.bfloat16)
    x = torch.randn(5, 4, 256).bfloat16()
    w = torch.randn(96, 256) / 16
    wq, ws = [], []
    # weight 32x32 block quant, like the checkpoint
    wb = w.view(3, 32, 8, 32)
    s = K._fast_round_scale(wb.abs().amax((1, 3)), 1 / 448)
    wq = (wb / s[:, None, :, None]).clamp(-448, 448).to(torch.float8_e4m3fn).view(96, 256)
    ws = K._to_scale_dtype(s, torch.float8_e8m0fnu)
    a, a_s = K.act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu)
    c = K.fp8_gemm(a, a_s, wq, ws, torch.float8_e8m0fnu, 32)
    wdeq = (wq.float().view(3, 32, 8, 32) * s[:, None, :, None]).view(96, 256)
    adeq = (a.float().unflatten(-1, (8, 32)) * a_s.float().unsqueeze(-1)).flatten(-2)
    ref = adeq @ wdeq.T
    assert c.dtype == torch.bfloat16 and c.shape == (5, 4, 96)
    assert rel(c, ref) < 4e-3
    assert rel(c, x.float() @ w.float().T) < 0.08
    torch.set_default_dtype(torch.float32)


def test_fp4_gemm():
    torch.set_default_dtype(torch.bfloat16)
    x = torch.randn(9, 128).bfloat16()
    w = torch.randn(64, 128) / 11
    s = K._fast_round_scale(w.view(64, 4, 32).abs().amax(-1), 1 / 6)
    codes = K._e2m1_codes((w.view(64, 4, 32) / s[..., None]).clamp(-6, 6)).view(64, 128)
    packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).to(torch.uint8).view(torch.float4_e2m1fn_x2)
    # decode exactly like convert.py cast_e2m1fn_to_e4m3fn (low nibble first, FP4_TABLE)
    FP4_TABLE = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])
    u = packed.view(torch.uint8)
    wdeq = torch.stack([FP4_TABLE[(u & 0xF).long()], FP4_TABLE[((u >> 4) & 0xF).long()]], -1).flatten(1)
    wdeq = wdeq * s.repeat_interleave(32, -1)
    a, a_s = K.act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu)
    c = K.fp4_gemm(a, a_s, packed, K._to_scale_dtype(s, torch.float8_e8m0fnu), torch.float8_e8m0fnu, 32)
    adeq = (a.float().view(9, 4, 32) * a_s.float().unsqueeze(-1)).flatten(-2)
    assert rel(c, adeq @ wdeq.T) < 4e-3
    assert rel(c, x.float() @ w.float().T) < 0.2
    torch.set_default_dtype(torch.float32)


def test_fp4_act_quant():
    x = torch.randn(3, 5, 128).bfloat16()
    for sd, bs in ((torch.float8_e8m0fnu, 32), (torch.float8_e4m3fn, 16)):
        y = x.clone()
        K.fp4_act_quant(y, bs, True, sd)
        packed, s = K.fp4_act_quant(x, bs, False, sd)
        deq = (K._unpack_fp4(packed).unflatten(-1, (-1, bs)) * s.float().unsqueeze(-1)).flatten(-2)
        assert torch.equal(deq.bfloat16(), y), sd
        assert rel(y, x) < 0.15
    # RNE ties on the e2m1 grid
    v = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, -0.25, -5.0, 5.01])
    assert K._e2m1_round(v).tolist() == [0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0, 0.0, -4.0, 6.0]


def test_sparse_attn():
    b, m, h, d, n, topk = 2, 5, 8, 512, 150, 140
    q = torch.randn(b, m, h, d).bfloat16()
    kv = torch.randn(b, n, d).bfloat16()
    sink = torch.randn(h) * 0.3
    idx = torch.stack([torch.randperm(n)[:topk] for _ in range(b * m)]).view(b, m, topk).int()
    idx[0, 0, 100:] = -1
    idx[1, 2, :] = -1  # a row with nothing to attend to -> all-zero output
    scale = d**-0.5
    o = K.sparse_attn(q, kv, sink, idx, scale)
    # dense fp32 reference with an explicit sink column
    kvg = kv[torch.arange(b).view(b, 1, 1), idx.long().clamp_min(0)].float()
    s = torch.einsum("bmhd,bmkd->bmhk", q.float(), kvg) * scale
    s = s.masked_fill((idx == -1).unsqueeze(2), -torch.inf)
    s = torch.cat([s, sink.view(1, 1, h, 1).expand(b, m, h, 1)], -1)
    p = torch.softmax(s, -1)[..., :-1]
    ref = torch.einsum("bmhk,bmkd->bmhd", p, kvg)
    assert rel(o, ref) < 1e-2, rel(o, ref)
    assert torch.all(o[1, 2] == 0)


def test_sinkhorn():
    mixes = torch.randn(2, 3, 24)
    pre, post, comb = K.hc_split_sinkhorn(mixes, torch.tensor([0.3, 0.3, 0.3]), torch.randn(24), 4, 20, 1e-6)
    assert pre.shape == (2, 3, 4) and post.shape == (2, 3, 4) and comb.shape == (2, 3, 4, 4)
    assert torch.allclose(comb.sum(-2), torch.ones(2, 3, 4), atol=1e-5)
    assert torch.allclose(comb.sum(-1), torch.ones(2, 3, 4), atol=1e-3)
    assert torch.all((post > 0) & (post < 2)) and torch.all(pre > 0)


if __name__ == "__main__":
    torch.manual_seed(0)
    for name, fn in list(globals().items()):
        if name.startswith("test_"):
            fn()
            print("ok", name)