from pathlib import Path
import pytest
import torch
from kernels import get_local_kernel

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

def inputs(rows,scale=1):
    torch.manual_seed(51)
    x=(torch.randn(rows,2048,device='cuda')*scale).bfloat16()
    w=torch.randn(2048,device='cuda').bfloat16()
    p=torch.randperm(2048,device='cuda').int()
    s=(torch.randint(0,2,(2048,),device='cuda')*2-1).to(torch.int8)
    bounds=torch.linspace(-.065,.065,15,device='cuda')
    return x,w,p,s,bounds

@pytest.mark.parametrize('rows',[0,1,2,8,32])
@pytest.mark.parametrize('scale',[0,1e-9,1,100])
def test_continuous_formula(rows,scale):
    x,w,p,s,b=inputs(rows,scale)
    packed,norms=kernel.rmsquant(x,w,p,s,b)
    y=x.float()*torch.rsqrt(x.float().square().mean(-1,keepdim=True)+1e-6)*w.float()
    expected_norm=y.norm(dim=-1)
    y=y/(expected_norm[:,None]+1e-8)
    y=y[:,p.long()]*s.float()
    width=1
    while width<2048:
        z=y.reshape(rows,2048//(2*width),2,width)
        a,c=z[:,:,0],z[:,:,1]
        y=torch.stack([a+c,a-c],dim=2).reshape(rows,2048)
        width*=2
    codes=torch.bucketize(y/(2048**.5),b)
    expected=(codes[:,0::2]|(codes[:,1::2]<<4)).byte()
    torch.testing.assert_close(norms,expected_norm,rtol=2e-6,atol=1e-9)
    # Only boundaries within FP32 reduction noise may differ.
    if rows:
        assert (expected!=packed).float().mean().item()<.001
    assert torch.isfinite(norms).all()

def test_graph_and_invalid_contract():
    args=inputs(2)
    expected=kernel.rmsquant(*args)
    stream=torch.cuda.Stream();stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        for _ in range(3):kernel.rmsquant(*args)
    torch.cuda.current_stream().wait_stream(stream)
    graph=torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):result=kernel.rmsquant(*args)
    graph.replay();torch.cuda.synchronize()
    for a,b in zip(expected,result):torch.testing.assert_close(a,b,rtol=0,atol=0)
    with pytest.raises(RuntimeError,match='BF16'):
        kernel.rmsquant(args[0].float(),*args[1:])
    with pytest.raises(RuntimeError,match='positive'):
        kernel.rmsquant(*args,rms_eps=0)
    with pytest.raises(RuntimeError,match='permutation'):
        kernel.rmsquant(args[0],args[1],args[2].long(),*args[3:])