import torch

from orbitquant_wan_a2.nibbles import pack_uint4
from orbitquant_wan_a2.reference import w4a4_linear_reference
from orbitquant_wan_a2.triton_w4a4 import PackedActivation, w4a4_linear_triton


def test_gemm_native_transposed_weight_layout_cpu_reference():
    torch.manual_seed(17)
    m, n, k = 5, 13, 64
    cb = torch.linspace(-0.2, 0.2, 16, dtype=torch.float32)
    ac = torch.randint(0, 16, (m, k), dtype=torch.uint8)
    wc = torch.randint(0, 16, (n, k), dtype=torch.uint8)
    ap = pack_uint4(ac)
    wp_row = pack_uint4(wc)
    wp_runtime = wp_row.transpose(0, 1).contiguous()
    a_scale = torch.rand(m) + 0.1
    w_scale = torch.rand(n) + 0.1
    bias = torch.randn(n, dtype=torch.bfloat16)
    packed_a = PackedActivation(ap, a_scale, k)
    got = w4a4_linear_triton(packed_a, wp_runtime, w_scale, cb, bias)
    ref = w4a4_linear_reference(ap, a_scale, wp_row, w_scale, cb, k, bias)
    assert torch.equal(got, ref)