"""Recompute stored results from serialized weights, independent of training helpers."""
from pathlib import Path
import json
import numpy as np

root=Path(__file__).parent
r=json.loads((root/'results.json').read_text())
d=np.load(root/'data.npz'); seq=d['test']
x=np.stack([seq[:,:-2],seq[:,1:-1]],axis=-1).reshape(-1,2)
y=seq[:,2:].ravel()
assert len(y)==4096
count=0
for run in r['runs']:
    assert run['selected_step']==min(run['history'],key=lambda a:a['dev_nll'])['step']
    for row in run['results']:
        folder=root/f"seed{run['seed']}_{row['name']}"
        manifest=json.loads((folder/'manifest.json').read_text());p={}
        for key,shape in r['shapes'].items():
            if key not in manifest:
                p[key]=np.fromfile(folder/f'{key}.f32',dtype='<f4').reshape(shape)
                continue
            a=manifest[key];scales=np.fromfile(folder/f'{key}.scale',dtype='<f4')
            raw=(folder/f'{key}.bin').read_bytes()
            if a['bits']==4:
                codes=np.array([v for b in raw for v in ((b & 15)-8,(b >> 4)-8)],dtype='float32')
                assert codes.min()>=-7 and codes.max()<=7
            else:
                codes=np.frombuffer(raw,dtype='int8').astype('float32')
            p[key]=(codes.reshape(-1,a['group'])*scales[:,None]).reshape(shape)
        h=np.tanh(p['E'][x].reshape(-1,16)@p['W1']+p['b1'])
        logits=(h@p['W2']+p['b2']).astype('float64')
        shifted=logits-logits.max(1,keepdims=True)
        nll=np.log(np.exp(shifted).sum(1))-shifted[np.arange(len(y)),y]
        np.testing.assert_allclose(nll,np.load(folder/'token_nll.npy'),rtol=0,atol=1e-12)
        assert abs(np.exp(nll.mean())-row['ppl'])<1e-12
        assert sum(f.stat().st_size for f in folder.iterdir() if f.suffix in ('.bin','.scale','.f32'))==row['payload_bytes']
        for label in ('cached_forward','decode_and_forward'):
            a=row[label];assert len(a['raw_us'])==101
            assert float(np.median(a['raw_us']))==a['median_us']
        count+=1
# Test all symmetric integer codes, including both extremes, with an independent decoder.
for bits in (4,8):
    maxq=2**(bits-1)-1
    q=np.arange(-maxq,maxq+1,dtype='int16')
    q=np.concatenate((q,[0]))
    if bits==4:
        u=(q+8).astype('uint8');b=(u[::2]|(u[1::2]<<4)).tobytes()
        decoded=np.array([v for c in b for v in ((c&15)-8,(c>>4)-8)])
    else:decoded=np.frombuffer(q.astype('int8').tobytes(),dtype='int8')
    np.testing.assert_array_equal(q,decoded)
print(f'PASS: {count} serialized model variants; 4096 token NLLs each; PPL, bytes, checkpoints, timing medians, all symmetric integer codes')