import json
from pathlib import Path
import pytest
import torch
from kernels import get_local_kernel
from orbitquant.kernels.native_packed_matmul import matmul_packed_w4a4_int8_with_native_kernel as original

kernel=get_local_kernel(Path(__file__).resolve().parents[1] / 'build', 'orbitquant_gemv')

@pytest.mark.parametrize('rows,n,k',[(1,2048,2048),(2,6144,2048),(8,2048,6144),(1,129,128),(2,1024,2048)])
@pytest.mark.parametrize('dtype',[torch.bfloat16,torch.float16])
def test_exact_integer_gemv(rows,n,k,dtype):
    torch.manual_seed(123)
    x=torch.randint(0,256,(rows,k//2),device='cuda',dtype=torch.uint8)
    w=torch.randint(0,256,(n*k//2,),device='cuda',dtype=torch.uint8)
    xn=torch.rand(rows,device='cuda');wn=torch.rand(n,device='cuda',dtype=torch.bfloat16)
    ac=torch.arange(-8,8,device='cuda',dtype=torch.int8);wc=ac.flip(0).contiguous()
    bias=None
    kw=dict(activation_scale=.03125,weight_scale=.0625,bias=bias,output_dtype=dtype)
    ref=original(x,w,xn,wn,ac,wc,out_features=n,in_features=k,**kw)
    out=kernel.gemv(x,w,xn,wn,ac,wc,**kw)
    torch.testing.assert_close(out,ref,rtol=0,atol=0)
    stream=torch.cuda.Stream();stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        for _ in range(3):kernel.gemv(x,w,xn,wn,ac,wc,**kw)
    torch.cuda.current_stream().wait_stream(stream)
    graph=torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph): captured=kernel.gemv(x,w,xn,wn,ac,wc,**kw)
    graph.replay();torch.cuda.synchronize()
    torch.testing.assert_close(captured,ref,rtol=0,atol=0)

def test_rejects_bad_weight_shape():
    x=torch.zeros(1,64,device='cuda',dtype=torch.uint8)
    with pytest.raises(RuntimeError,match='packed weight size'):
        kernel.gemv(x,x.flatten(),torch.ones(1,device='cuda'),torch.ones(2,device='cuda',dtype=torch.bfloat16),torch.zeros(16,device='cuda',dtype=torch.int8),torch.zeros(16,device='cuda',dtype=torch.int8),activation_scale=1,weight_scale=1)