"""Exact attention gate laws for all finite BF16/F16 values and dense F32 inputs."""
import hashlib,json,subprocess,sys
from pathlib import Path
checkout,out=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.attention_math import attention_gate
from mtplx.attention_context import attention_phase
from mtplx.models.qwen4_exp import _verify_inject
out.mkdir(parents=True,exist_ok=True)
for name,dtype in [('bf16',mx.bfloat16),('f16',mx.float16),('f32',mx.float32)]:
    if name=='bf16':
        values=(np.arange(65536,dtype=np.uint32)<<16).view(np.float32)
    elif name=='f16':
        values=np.arange(65536,dtype=np.uint16).view(np.float16).astype(np.float32)
    else:
        values=np.concatenate([np.linspace(-30,30,32768,dtype=np.float32),np.random.default_rng(5).normal(0,4,32768).astype(np.float32)])
    values=values[np.isfinite(values)]
    gate=mx.array(values).astype(dtype)
    output=mx.array(((np.arange(values.size,dtype=np.int64)*17)%97-48).astype(np.float32)/32).astype(dtype)
    mx.eval(gate,output)
    for part,a in [('gate',gate),('output',output)]:
        (out/f'{name}-{part}.f32').write_bytes(np.asarray(a.astype(mx.float32)).tobytes())
    for phase in ['prefill','ar_decode','decode_verify']:
        with attention_phase(phase): result=attention_gate(output,gate)
        raw=np.asarray(result.astype(mx.float32)).tobytes()
        (out/f'{name}-{phase}.f32').write_bytes(raw)
        print(json.dumps(dict(case='gate',dtype=name,phase=phase,count=values.size,sha256=hashlib.sha256(raw).hexdigest())),flush=True)
        if name=='bf16':
            inject=_verify_inject(gate,4) if phase=='decode_verify' else 2.0*mx.sigmoid(gate/4)
            raw=np.asarray(inject.astype(mx.float32)).tobytes()
            (out/f'inject-{phase}.f32').write_bytes(raw)
            print(json.dumps(dict(case='inject',dtype=name,phase=phase,count=values.size,sha256=hashlib.sha256(raw).hexdigest())),flush=True)