import torch
from orbitquant_wan_a2.nibbles import pack_uint4
from orbitquant_wan_a2.reference import w4a4_linear_reference


def test_reference_matches_explicit_bf16_dequant():
    torch.manual_seed(3)
    m, n, k = 7, 11, 64
    cb = torch.linspace(-0.2, 0.2, 16)
    ac = torch.randint(0,16,(m,k),dtype=torch.uint8)
    wc = torch.randint(0,16,(n,k),dtype=torch.uint8)
    a_s = torch.rand(m) + 0.1
    w_s = torch.rand(n) + 0.1
    bias = torch.randn(n, dtype=torch.bfloat16)
    out = w4a4_linear_reference(pack_uint4(ac), a_s, pack_uint4(wc), w_s, cb, k, bias)
    a = (cb[ac.long()] * a_s[:,None]).to(torch.bfloat16)
    w = (cb[wc.long()] * w_s[:,None]).to(torch.bfloat16)
    ref = (a.float() @ w.float().T + bias.float()).to(torch.bfloat16)
    assert torch.equal(out, ref)