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

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

def quantize(x):
 scale=x.float().abs().amax(-1).clamp_min(1e-12)/127
 return (x.float()/scale[...,None]).round().clamp(-127,127).to(torch.int8),scale

def reference(q,k,v,ks,vs,lengths):
 out=[]
 for b,length in enumerate(lengths):
  if not length:out.append(torch.zeros_like(q[b],dtype=torch.float32));continue
  key=(k[b,:length].float()*ks[b,:length,None].movedim(1,2)).repeat_interleave(2,1)
  value=(v[b,:length].float()*vs[b,:length,None].movedim(1,2)).repeat_interleave(2,1)
  scores=torch.einsum('hd,thd->ht',q[b].float(),key)*(128**-.5)
  out.append(torch.einsum('ht,thd->hd',scores.softmax(-1),value))
 return torch.stack(out)

@pytest.mark.parametrize('parts',[16,32,64,128])
@pytest.mark.parametrize('dtype',[torch.bfloat16,torch.float16])
@pytest.mark.parametrize('capacity,lengths',[(129,[0,129]),(513,[513,129]),(11053,[2053,6053])])
def test_quantized_lengths_and_graph(parts,dtype,capacity,lengths):
 torch.manual_seed(995)
 q=torch.randn(2,16,128,device='cuda',dtype=dtype)
 k,ks=quantize(torch.randn(2,capacity,8,128,device='cuda',dtype=dtype))
 v,vs=quantize(torch.randn(2,capacity,8,128,device='cuda',dtype=dtype))
 for b,length in enumerate(lengths):ks[b,length:]=float('nan');vs[b,length:]=float('nan')
 used=torch.tensor(lengths,device='cuda',dtype=torch.int32)
 out=kernel.attention(q,k,v,ks,vs,used,parts)
 assert torch.isfinite(out).all()
 torch.testing.assert_close(out.float(),reference(q,k,v,ks,vs,lengths),atol=.001,rtol=.01)
 stream=torch.cuda.Stream();stream.wait_stream(torch.cuda.current_stream())
 with torch.cuda.stream(stream):
  for _ in range(3):kernel.attention(q,k,v,ks,vs,used,parts)
 torch.cuda.current_stream().wait_stream(stream)
 graph=torch.cuda.CUDAGraph()
 with torch.cuda.graph(graph):captured=kernel.attention(q,k,v,ks,vs,used,parts)
 short=[min(x,7) for x in lengths];used.copy_(torch.tensor(short,device='cuda',dtype=torch.int32));graph.replay();torch.cuda.synchronize()
 torch.testing.assert_close(captured.float(),reference(q,k,v,ks,vs,short),atol=.001,rtol=.01)

def test_invalid_scale_and_cache_contracts():
 q=torch.zeros(1,16,128,device='cuda',dtype=torch.bfloat16)
 k=torch.zeros(1,10,8,128,device='cuda',dtype=torch.int8)
 s=torch.ones(1,10,8,device='cuda');used=torch.ones(1,device='cuda',dtype=torch.int32)
 with pytest.raises(RuntimeError,match='scales'):kernel.attention(q,k,k,s.half(),s,used)
 with pytest.raises(RuntimeError,match='scales'):kernel.attention(q,k,k,s[:,:,:7],s,used)
 with pytest.raises(RuntimeError,match='KV dtype'):kernel.attention(q,k.bfloat16(),k,s,s,used)


def test_unaligned_cache_view():
 torch.manual_seed(109)
 q=torch.randn(1,16,128,device='cuda',dtype=torch.bfloat16)
 shape=(1,129,8,128);n=129*8*128
 backing=torch.randint(-127,128,(n+1,),device='cuda',dtype=torch.int8)
 k=backing[1:].view(shape);v=k.clone();scales=torch.full((1,129,8),.02,device='cuda');used=torch.tensor([127],device='cuda',dtype=torch.int32)
 out=kernel.attention(q,k,v,scales,scales,used,32)
 torch.testing.assert_close(out.float(),reference(q,k,v,scales,scales,[127]),atol=.001,rtol=.01)


@pytest.mark.parametrize('dtype',[torch.bfloat16,torch.float16])
@pytest.mark.parametrize('batch',[1,2])
@pytest.mark.parametrize('parts',[32,64,128])
@pytest.mark.parametrize('length',[129,6053])
def test_original_reduction_order(dtype,batch,parts,length):
 control=get_local_kernel(Path(__file__).resolve().parents[2]/'int8-kv/build','int8_kv_runtime')
 torch.manual_seed(110)
 q=torch.randn(batch,16,128,device='cuda',dtype=dtype)
 k,ks=quantize(torch.randn(batch,length+37,8,128,device='cuda',dtype=dtype))
 v,vs=quantize(torch.randn(batch,length+37,8,128,device='cuda',dtype=dtype))
 used=torch.tensor([length-i*31 for i in range(batch)],device='cuda',dtype=torch.int32)
 torch.testing.assert_close(kernel.attention(q,k,v,ks,vs,used,parts),control.attention(q,k,v,ks,vs,used,parts),atol=0,rtol=0)