"""Isolate the extra oracle snapshot that changed the reported key's zero sign."""
import hashlib,json,subprocess,sys
from pathlib import Path
checkout,native,old_oracle,out=map(Path,sys.argv[1:5])
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.models.qwen4_exp import QSACache
from mtplx.cache_state import snapshot_cache,snapshot_cache_lazy_hybrid
assert mx.__version__=='0.32.2'
def data(a):return np.asarray(a.view(mx.uint16 if a.dtype==mx.bfloat16 else mx.uint32)).tobytes()
def cache(a):
 c=QSACache();c.state=(a,a,None,None);return c
cases=[]
for dtype in [mx.bfloat16,mx.float32]:
 bits=np.array([0x80000000,0,0xbf800000,0x3f800000,0,0x80000000,0x40000000,0xc0000000],dtype=np.uint32)
 a=mx.array(bits.view(np.float32).reshape(1,2,1,4)).astype(dtype)
 c=cache(a);raw=c.state[0];lazy=snapshot_cache_lazy_hybrid([c]);deep=snapshot_cache([c])
 row=dict(dtype=str(dtype),shape=list(a.shape),raw=data(raw).hex(),lazy=data(lazy.states[0][0]).hex(),deep=data(deep.states[0][0]).hex())
 assert row['raw']==row['lazy'] and row['raw']!=row['deep']
 cases.append(row)
traces=[]
for phase in ['cold-prefill','cold','warm']:
 f=native/(phase+'-key35.f32');before=f.read_bytes();shape=(1,2,len(before)//(2*256*4),256)
 a=mx.array(np.frombuffer(before,dtype='<f4').copy().reshape(shape)).astype(mx.bfloat16)
 def f32(x):return np.asarray(x.astype(mx.float32)).astype('<f4').tobytes()
 assert f32(a)==before,'input reconstruction must retain every bit'
 c=cache(a);raw=f32(c.state[0]);lazy=f32(snapshot_cache_lazy_hybrid([c]).states[0][0]);deep=f32(snapshot_cache([c]).states[0][0])
 original=(old_oracle/(phase+'-key35.f32')).read_bytes()
 assert raw==lazy==before
 assert deep==original,'extra clone must reproduce the original mismatch exactly'
 x=np.frombuffer(raw,np.uint32).reshape(shape);y=np.frombuffer(deep,np.uint32).reshape(shape)
 coords=np.argwhere(x!=y).tolist()
 for at in coords:assert x[tuple(at)]==0x80000000 and y[tuple(at)]==0
 traces.append(dict(phase=phase,shape=shape,coordinates=coords,raw_sha256=hashlib.sha256(raw).hexdigest(),lazy_sha256=hashlib.sha256(lazy).hexdigest(),deep_sha256=hashlib.sha256(deep).hexdigest(),original_oracle_sha256=hashlib.sha256(original).hexdigest()))
out.write_text(json.dumps(dict(revision='9882703f3105363ddc37eca9f97aa09a1d387112',mlx='0.32.2',cases=cases,real_key_traces=traces),indent=2)+'\n')
print(json.dumps(dict(cases=len(cases),real_key_traces=traces)),flush=True)