import sys, math, json, random
import numpy as np
import torch
import torch.nn as nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model
from bench_imm import IMMMonitor, seed_all, vec_norm, EPOCHS, BATCH

def run(seed, delay=0):
    seed_all(seed); ds=get_dataset('tabular',seed,n_train=400,n_test=400)
    net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']); opt=torch.optim.Adam(net.parameters(),lr=.01)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net.to(dev); x,y=ds['xtr'].to(dev),ds['ytr'].to(dev); mon=IMMMonitor(); queued=[]; p=[]; alarms=[]; inn=[]
        for ep in range(8):
            net.train(); perm=torch.randperm(len(x),device=dev)
            for j in range(0,len(x),BATCH):
                idx=perm[j:j+BATCH]; loss=nn.MSELoss()(net(x[idx]),y[idx]); opt.zero_grad(); loss.backward()
                grads=[None if q.grad is None else q.grad.detach().clone() for q in net.parameters()]
                queued.append(grads)
                use=queued[max(0,len(queued)-1-delay)]
                for q,g in zip(net.parameters(),use):
                    if g is not None: q.grad=g
                gn=math.sqrt(sum(float((g**2).sum()) for g in use if g is not None)); pn=vec_norm(net)
                z=np.log1p(np.array([float(loss),float(loss),gn,pn,.01*gn]))
                al=mon.step(z); alarms.append(al); p.append(mon.pi.copy()); inn.append(mon.last['innovation'])
                opt.step()
        p=np.asarray(p); return {'delay':delay,'final_posterior':p[-1].tolist(),'min_no_delay_posterior':float(p[:,0].min()),'alarm_count':int(sum(alarms)),'first_alarm':next((i+1 for i,a in enumerate(alarms) if a),None),'mean_innovation':float(np.mean(inn))}
    except RuntimeError:
        return {'delay':delay,'error':'runtime failure'}

def main():
    clean=[run(s,0) for s in range(8)]; stale=[run(s,2) for s in range(8)]
    out={'clean':clean,'stale_delay_2':stale,'summary':{'clean_alarm_rate':float(np.mean([x['alarm_count']>0 for x in clean])),'stale_alarm_rate':float(np.mean([x['alarm_count']>0 for x in stale])),'stale_first_alarm_median':float(np.median([x['first_alarm'] for x in stale if x['first_alarm'] is not None])) if any(x['first_alarm'] is not None for x in stale) else None}}
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()