"""A deterministic CPU training replay lab inspired by NVIDIA Megatron Core.
Run: OPENBLAS_NUM_THREADS=1 python experiment.py
Dependencies: Python, NumPy, Matplotlib. No downloads, network, model or GPU.
"""
import os
os.environ.setdefault('MPLCONFIGDIR','/tmp/engineering-blog-mpl')
import copy,json,hashlib,csv
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
ROOT=Path(__file__).resolve().parent if '__file__' in globals() else Path.cwd()
STEPS=100;CUT=40;BATCH=64;D=24
COLORS={'Same seed + full state':'#eeeeee','Full checkpoint':'#eeeeee','Reversed reduction':'#333333','Optimizer reset':'#9a9a9a','Sampler RNG reset':'#536b7a'}

def digest(a):
    a=np.ascontiguousarray(a)
    return hashlib.sha256(str(a.shape).encode()+str(a.dtype).encode()+a.tobytes()).hexdigest()

def make_data():
    r=np.random.default_rng(20261106)
    x=r.normal(size=(512,D)).astype(np.float32)
    truth=r.normal(size=D).astype(np.float32)
    y=(x@truth+.2*np.sin(3*x[:,0])+.03*r.normal(size=512)).astype(np.float32)
    return x,y

def init_state():
    return {'w':np.random.default_rng(17).normal(0,.01,size=D).astype(np.float32),'m':np.zeros(D,np.float32),'v':np.zeros(D,np.float32),'step':0,'rng':copy.deepcopy(np.random.default_rng(99).bit_generator.state)}

def state_digest(s):
    return hashlib.sha256((''.join(digest(s[k]) for k in ['w','m','v'])+json.dumps({'step':s['step'],'rng':s['rng']},sort_keys=True)).encode()).hexdigest()

def reduce_rows(x,reverse=False):
    out=np.zeros(x.shape[1],np.float32)
    for row in (x[::-1] if reverse else x):
        out=np.add(out,row,dtype=np.float32)
    return out

def step(state,data,reverse=False):
    x,y=data
    rng=np.random.default_rng();rng.bit_generator.state=copy.deepcopy(state['rng'])
    idx=rng.choice(len(x),BATCH,replace=False)
    state['rng']=copy.deepcopy(rng.bit_generator.state)
    xb=x[idx];yb=y[idx]
    pred=np.sum(xb*state['w'],axis=1,dtype=np.float32)
    residual=pred-yb
    per_example=(np.float32(2/BATCH)*residual[:,None]*xb).astype(np.float32)
    g=reduce_rows(per_example,reverse)
    input_hash=digest(per_example);grad_hash=digest(g)
    state['step']+=1;t=state['step']
    b1=np.float32(.9);b2=np.float32(.999);lr=np.float32(.04)
    state['m']=(b1*state['m']+(np.float32(1)-b1)*g).astype(np.float32)
    state['v']=(b2*state['v']+(np.float32(1)-b2)*g*g).astype(np.float32)
    mh=state['m']/np.float32(1-float(b1)**t)
    vh=state['v']/np.float32(1-float(b2)**t)
    state['w']=(state['w']-lr*mh/(np.sqrt(vh)+np.float32(1e-8))).astype(np.float32)
    # Float64 evaluation avoids hiding loss changes inside its own float32 reduction.
    residual_eval=x.astype(np.float64)@state['w'].astype(np.float64)-y
    loss=float(np.mean(residual_eval**2))
    return {'step':t,'loss':loss,'printed_loss':f'{loss:.4f}','weights':state['w'].copy(),'weight_hash':digest(state['w']),'state_hash':state_digest(state),'gradient_input_hash':input_hash,'gradient_hash':grad_hash,'sample_hash':digest(idx)}

def run(data,state=None,steps=STEPS,reverse=False):
    s=init_state() if state is None else copy.deepcopy(state);trace=[]
    for _ in range(steps):trace.append(step(s,data,reverse))
    return s,trace

def serial_trace(trace):
    return [{k:v for k,v in t.items() if k!='weights'} for t in trace]

def run_experiment():
    data=make_data();final,ref=run(data)
    _,same=run(data)
    checkpoint,prefix=run(data,steps=CUT)
    restored,tail=run(data,checkpoint,STEPS-CUT)
    _,reverse=run(data,reverse=True)
    broken_optimizer=copy.deepcopy(checkpoint);broken_optimizer['m'][:]=0;broken_optimizer['v'][:]=0
    # Model the common partial-restart behavior: reset both Adam moments and its step counter.
    broken_optimizer['step']=0
    _,badopt=run(data,broken_optimizer,STEPS-CUT)
    # Report training positions on the common x-axis; the optimizer counter remains reset internally.
    for j,t in enumerate(badopt):t['step']=CUT+j+1
    broken_rng=copy.deepcopy(checkpoint);broken_rng['rng']=copy.deepcopy(np.random.default_rng(99).bit_generator.state)
    _,badrng=run(data,broken_rng,STEPS-CUT)
    traces={'Reference':ref,'Same seed + full state':same,'Full checkpoint':prefix+tail,'Reversed reduction':reverse,'Optimizer reset':prefix+badopt,'Sampler RNG reset':prefix+badrng}
    rows=[]
    for name,trace in traces.items():
        if name=='Reference':continue
        for a,b in zip(ref,trace):
            rows.append({'scenario':name,'step':a['step'],'exact_weights_match':a['weight_hash']==b['weight_hash'],'exact_state_match':a['state_hash']==b['state_hash'],'same_4dp_loss':a['printed_loss']==b['printed_loss'],'max_weight_error':float(np.max(np.abs(a['weights'].astype(float)-b['weights']))),'loss':b['loss'],'reference_loss':a['loss']})
    return data,traces,rows,checkpoint,final,restored

def validate(data,traces,rows,checkpoint,final,restored):
    tests=[]
    def check(name,ok):assert ok,name;tests.append(name)
    check('same-state repeat matches at every step',all(r['exact_state_match'] for r in rows if r['scenario']=='Same seed + full state'))
    check('full checkpoint matches at every step',all(r['exact_state_match'] for r in rows if r['scenario']=='Full checkpoint'))
    check('full resumed final arrays are byte equal',all(final[k].tobytes()==restored[k].tobytes() for k in ['w','m','v']))
    check('full resumed RNG and step equal',final['step']==restored['step'] and final['rng']==restored['rng'])
    a=traces['Reference'][0];b=traces['Reversed reduction'][0]
    check('reduction comparison has identical ordered inputs',a['gradient_input_hash']==b['gradient_input_hash'])
    check('reduction order changes first gradient bytes',a['gradient_hash']!=b['gradient_hash'])
    check('changed order is hidden by rounded logging somewhere',any(not r['exact_weights_match'] and r['same_4dp_loss'] for r in rows if r['scenario']=='Reversed reduction'))
    for name in ['Optimizer reset','Sampler RNG reset']:
        check(name+' matches prefix',all(r['exact_state_match'] for r in rows if r['scenario']==name and r['step']<=CUT))
        check(name+' breaks replay after restore',any(not r['exact_state_match'] for r in rows if r['scenario']==name and r['step']>CUT))
    check('training loss decreases',traces['Reference'][-1]['loss']<traces['Reference'][0]['loss'])
    check('fingerprint detects permutation',digest(np.array([1,2,3]))!=digest(np.array([3,2,1])))
    # Different dtypes and shapes must not accidentally share a byte-only digest.
    check('fingerprint includes shape and dtype',digest(np.array([1,2],np.int32))!=digest(np.array([[1,2]],np.int32)))
    return tests

def summarize(rows,traces):
    out=[]
    for name in COLORS:
        r=[r for r in rows if r['scenario']==name]
        bad=[v['step'] for v in r if not v['exact_weights_match']]
        out.append({'scenario':name,'first_weight_mismatch':min(bad) if bad else None,'weight_mismatch_steps':len(bad),'hidden_weight_mismatch_steps':sum(not v['exact_weights_match'] and v['same_4dp_loss'] for v in r),'maximum_weight_difference':max(v['max_weight_error'] for v in r),'final_loss':r[-1]['loss']})
    return out

def make_figures(rows,traces,out=ROOT):
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':11,'axes.spines.top':False,'axes.spines.right':False,'axes.edgecolor':'#555555','axes.labelcolor':'#333333','text.color':'#222222','xtick.color':'#555555','ytick.color':'#555555','axes.titleweight':'normal','axes.titlesize':12,'svg.fonttype':'none','figure.facecolor':'white','axes.facecolor':'white','grid.color':'#dddddd','grid.linewidth':.6})
    fig,axs=plt.subplots(2,1,figsize=(9.2,6.4),gridspec_kw={'height_ratios':[1.1,1.5]})
    names=list(COLORS);matrix=np.array([[int(r['exact_weights_match']) for r in rows if r['scenario']==n] for n in names])
    from matplotlib.colors import ListedColormap
    from matplotlib.patches import Patch
    axs[0].imshow(matrix,aspect='auto',extent=(.5,100.5,4.5,-.5),cmap=ListedColormap(['#536b7a','#eeeeee']),vmin=0,vmax=1,interpolation='nearest')
    axs[0].set_yticks(range(5),names);axs[0].set_xticks([1,20,40,60,80,100]);axs[0].tick_params(axis='y',length=0);axs[0].set_title('Exact weight agreement with the reference',loc='left',pad=30);axs[0].axvline(CUT+.5,color='#333333',lw=.8)
    axs[0].legend(handles=[Patch(facecolor='#eeeeee',label='Matching bytes'),Patch(facecolor='#536b7a',label='Different bytes')],loc='lower left',bbox_to_anchor=(0,1.01),ncol=2,frameon=False,fontsize=9,borderaxespad=0)
    styles={'Reversed reduction':('#333333','-'),'Optimizer reset':('#9a9a9a','--'),'Sampler RNG reset':('#536b7a',':')}
    for n,(c,ls) in styles.items():
        aa=[r for r in rows if r['scenario']==n and r['max_weight_error']>0]
        axs[1].plot([r['step'] for r in aa],[r['max_weight_error'] for r in aa],color=c,ls=ls,label=n,lw=1.6)
    axs[1].set_yscale('log');axs[1].set_xlabel('Training step');axs[1].set_ylabel('Maximum weight difference');axs[1].axvline(CUT+.5,color='#bbbbbb',ls=':',lw=1);axs[1].grid(axis='y');axs[1].legend(frameon=False,fontsize=9,ncol=3,loc='upper center',bbox_to_anchor=(.5,-.22))
    fig.text(.03,.012,'Float32 CPU study · checkpoint after step 40 · zero differences omitted from the logarithmic plot',fontsize=9,color='#666666')
    fig.tight_layout(rect=(0,.055,1,1),h_pad=2.5)
    for name in ['replay','cover']:
        fig.savefig(out/(name+'.svg'),bbox_inches='tight');fig.savefig(out/(name+'.png'),dpi=180,bbox_inches='tight')
    plt.close(fig)
    ref=traces['Reference'];rev=traces['Reversed reduction']
    fig,ax=plt.subplots(figsize=(7,4))
    ax.plot(range(1,STEPS+1),[x['loss'] for x in ref],color='#aaaaaa',lw=2.2,label='Reference');ax.plot(range(1,STEPS+1),[x['loss'] for x in rev],color='#333333',lw=1.2,ls='--',label='Reversed reduction')
    ax.set_xlabel('Training step');ax.set_ylabel('Evaluation mean squared error');ax.legend(frameon=False,fontsize=10);ax.grid(axis='y');fig.tight_layout();fig.savefig(out/'hidden-drift.svg',bbox_inches='tight');fig.savefig(out/'hidden-drift.png',dpi=180,bbox_inches='tight');plt.close(fig)

def main():
    data,traces,rows,checkpoint,final,restored=run_experiment();tests=validate(data,traces,rows,checkpoint,final,restored);summary=summarize(rows,traces)
    (ROOT/'results.json').write_text(json.dumps({'config':{'steps':STEPS,'checkpoint_step':CUT,'batch':BATCH,'dimensions':D,'rows':512,'data_seed':20261106,'weights_seed':17,'sampler_seed':99,'dtype':'float32','optimizer':'Adam','learning_rate':.04},'summary':summary,'tests':tests,'traces':{k:serial_trace(v) for k,v in traces.items()}},indent=2))
    with (ROOT/'comparisons.csv').open('w',newline='') as f:w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    np.savez(ROOT/'checkpoint.npz',w=checkpoint['w'],m=checkpoint['m'],v=checkpoint['v'])
    (ROOT/'checkpoint-state.json').write_text(json.dumps({'step':checkpoint['step'],'rng':checkpoint['rng'],'data_sha256':digest(data[0])+':'+digest(data[1]),'note':'Load checkpoint.npz with allow_pickle=False. Recreate the documented fixed data and configuration.'},indent=2))
    # Verify the file round-trip independently of the in-memory replay.
    loaded=np.load(ROOT/'checkpoint.npz',allow_pickle=False);meta=json.loads((ROOT/'checkpoint-state.json').read_text());s={k:loaded[k] for k in ['w','m','v']};s.update(step=meta['step'],rng=meta['rng']);disk_final,_=run(data,s,STEPS-CUT)
    assert state_digest(disk_final)==state_digest(final);tests.append('saved checkpoint file round-trip matches reference')
    make_figures(rows,traces)
    saved=json.loads((ROOT/'results.json').read_text());saved['tests']=tests;(ROOT/'results.json').write_text(json.dumps(saved,indent=2))
    (ROOT/'checks.json').write_text(json.dumps({'passed':tests,'count':len(tests)},indent=2))
    print(json.dumps({'checks_passed':len(tests),'summary':summary,'reference_final_loss':traces['Reference'][-1]['loss']},indent=2))

if __name__=='__main__':main()