import sys, copy, json, numpy as np, torch
import torch.nn as nn
sys.path.insert(0,'/home/maxwelhelp/all/math2nn')
import bench
from stage2_bench import setup, make_blocks, choose_order, EPOCHS, BATCH

# Retest the local prediction on one actually trained benchmark model.
def main():
    seed=0; cfg={'lr':0.0015,'weight_decay':0.0}
    ds, model=setup(seed)
    net, metric, hist=bench.train_model(model,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=0.0,log=lambda *a,**k:None)
    dev=next(net.parameters()).device
    x,y=ds['xtr'].to(dev),ds['ytr'].to(dev)
    loss_fn=lambda: nn.MSELoss()(net(x),y)
    bs=make_blocks(net)
    order,pred_rho,sens,mat=choose_order(net,loss_fn,bs,cfg['lr'])
    # Apply the same sequential gradient map to two nearby trained states.
    q=copy.deepcopy(net)
    qbs=make_blocks(q); q.to(dev)
    delta=[]
    with torch.no_grad():
        for p in q.parameters():
            z=torch.randn_like(p); p.add_(1e-4*z/max(float(z.norm()),1.0))
    def sweep(m, groups):
        for b in order:
            m.zero_grad(set_to_none=True)
            l=nn.MSELoss()(m(x),y); l.backward()
            with torch.no_grad():
                for p in groups[b]:
                    if p.grad is not None: p.add_(-cfg['lr']*p.grad)
    before=torch.cat([p.detach().reshape(-1) for p in net.parameters()])
    before_q=torch.cat([p.detach().reshape(-1) for p in q.parameters()])
    sweep(net,bs); sweep(q,qbs)
    after=torch.cat([p.detach().reshape(-1) for p in net.parameters()])
    after_q=torch.cat([p.detach().reshape(-1) for p in q.parameters()])
    observed=float((after_q-after).norm()/(before_q-before).norm())
    sig={'predicted_rho':float(pred_rho),'observed_local_perturbation_ratio':observed,
         'relative_gap_pct':float(abs(observed-pred_rho)/max(abs(pred_rho),1e-12)*100),
         'selected_order':order,'trained_test_mse':float(metric),
         'confirmed':bool(abs(observed-pred_rho)/max(abs(pred_rho),1e-12)<=0.20)}
    with open('bench_report.json') as f: rep=json.load(f)
    rep['mechanism_signature']=sig
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(sig,indent=2))
if __name__=='__main__': main()