"""Pinned row-map kernel, guarded stock chain and native synthetic bit receipts."""
import hashlib,json,subprocess,sys
from pathlib import Path
checkout,native=map(Path,sys.argv[1:3])
assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
sys.path.insert(0,str(checkout))
import mlx.core as mx
import numpy as np
from mtplx.kernels import moe_sorted_gather_nax as kernel
assert mx.__version__=='0.32.2'
def pattern(count,salt,scale=1):
 a=((np.arange(count,dtype=np.uint32)*17+salt*13)%257).astype(np.int32)-128
 return mx.array(a.astype(np.float32)/128*scale).astype(mx.bfloat16)
def pack(e,n,k,bits,group,salt,dtype):
 words=np.arange(e*n*k//32*bits,dtype=np.uint32)*np.uint32(2654435761)+np.uint32(12345)
 return (mx.array(words.reshape(e,n,k*bits//32)),pattern(e*n*k//group,salt,1/64).reshape(e,n,k//group).astype(dtype),pattern(e*n*k//group,salt+1,1/64).reshape(e,n,k//group).astype(dtype))
def digest(a):
 mx.eval(a)
 return hashlib.sha256(np.asarray(a.view(mx.uint16)).tobytes()).hexdigest()
count=0
for line in native.read_text().splitlines():
 if '"case":"synthetic"' not in line:continue
 case=json.loads(line[line.index('{'):])
 count+=1
 dtype={'BF16':mx.bfloat16,'F16':mx.float16}[case['dtype']]
 bits,group,rows=case['bits'],case['group'],case['rows']
 gate=pack(128,256,256,bits,group,7,dtype)
 up=pack(128,256,256,bits,group,9,dtype) if case['separate'] else None
 x=pattern(512*256,8).reshape(512,1,256).astype(dtype)
 r=np.arange(rows,dtype=np.uint32)
 ids=mx.array((r*125//rows)//2*2+1)
 row_map=mx.array((r*13)%512)
 assert kernel.applies(x,row_map,*gate,ids,group_size=group,bits=bits,mode='affine',up=up)
 kwargs=dict(group_size=group,bits=bits,swiglu=True,up=up)
 before=kernel.stats()
 y=kernel.gather_rows_qmm(x,row_map,*gate,ids,**kwargs)
 assert y is not None,kernel.stats()
 stock=kernel._stock(x,row_map,*gate,ids,**kwargs)
 actual=digest(y)
 assert actual==digest(stock)==case['sha256'],(case,actual,digest(stock))
 print(json.dumps(dict(case='synthetic_oracle',native=case,exact=True,kernel_before=before,kernel_after=kernel.stats())),flush=True)
 del x,row_map,ids,gate,up,y,stock
 mx.clear_cache()

assert count==72,count