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', 'yue2_qkv_fused')

@pytest.mark.parametrize('rows,n,k',[(1,2048,2048),(2,6144,2048),(8,2048,6144),(1,129,128),(2,1024,2048),(1,129,130),(1,129,16)])
@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.randint(-127,128,(16,),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)
    if k % 128 == 0:
        ref=original(x,w,xn,wn,ac,wc,out_features=n,in_features=k,**kw)
    else:
        def unpack(z):return torch.stack([z & 15,z >> 4],dim=-1).reshape(z.shape[0],-1).long().cpu()
        aa=ac.cpu().long()[unpack(x)];ww=wc.cpu().long()[unpack(w.reshape(n,k//2))]
        sums=(aa @ ww.T).to(device='cuda',dtype=torch.float32)
        ref=(sums*(xn[:,None]*wn.float()[None,:]*(kw['activation_scale']*kw['weight_scale']))).to(dtype)
    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)

@pytest.mark.parametrize('offset',[1,2,3])
def test_misaligned_contiguous_packed_inputs(offset):
    rows,n,k=1,129,128
    torch.manual_seed(10)
    x=torch.randint(0,256,(rows*k//2+offset,),device='cuda',dtype=torch.uint8)[offset:].view(rows,k//2)
    w=torch.randint(0,256,(n*k//2+offset,),device='cuda',dtype=torch.uint8)[offset:]
    xn=torch.rand(rows,device='cuda');wn=torch.rand(n,device='cuda',dtype=torch.bfloat16)
    ac=torch.randint(-127,128,(16,),device='cuda',dtype=torch.int8);wc=ac.flip(0).contiguous()
    kw=dict(activation_scale=.03125,weight_scale=.0625,bias=None,output_dtype=torch.bfloat16)
    ref=original(x.contiguous().clone(),w.contiguous().clone(),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)