import sys, json, copy, 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, evaluate, sweep_baseline, make_report

SEEDS=tuple(range(8)); SWEEP=tuple(range(4)); EPOCHS=18; BATCH=128
# Union of all learning rates: baseline and certified wrapper both see these.
GRID=[{'lr':1e-3,'M':1},{'lr':3e-3,'M':1},{'lr':1e-2,'M':1}]
IDEA_GRID=GRID

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)

def device_model(ds):
    # train_model's fallback is not usable because this is a modified optimizer loop.
    try:
        dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        return make_model('rnn_small', ds['input_shape'], ds['out_dim']).to(dev),dev
    except Exception:
        return make_model('rnn_small', ds['input_shape'], ds['out_dim']).to('cpu'),torch.device('cpu')

def hidden_energy(net,x):
    got={}
    def hook(mod, inp, out):
        h=out[1]
        got['v']=(h*h).mean()
    h=net.rnn.register_forward_hook(hook)
    try: net(x)
    finally: h.remove()
    return got.get('v', torch.tensor(0.,device=x.device))

def train(seed,cfg,certified=False, collect=False):
    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=160)
    try: net,dev=device_model(ds)
    except Exception: net=make_model('rnn_small',ds['input_shape'],ds['out_dim']); dev=torch.device('cpu')
    xtr,ytr=[ds[k].to(dev) for k in ('xtr','ytr')]; xte,yte=[ds[k].to(dev) for k in ('xte','yte')]
    lossf=nn.MSELoss(); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
    M=cfg.get('M',1); lam=0.02; c=1.0
    accepted=rejected=checks=violations=0; cert_margins=[]; prev_batch=None
    for ep in range(EPOCHS):
        net.train(); perm=torch.randperm(len(xtr),device=dev); opt.zero_grad(set_to_none=True)
        for bi,i in enumerate(range(0,len(xtr),BATCH)):
            ix=perm[i:i+BATCH]; xb,yb=xtr[ix],ytr[ix]
            loss=lossf(net(xb),yb); loss.backward()
            if ((bi+1)%M and i+BATCH<len(xtr)): continue
            # Save proposed state and old energy on a fixed current probe.
            if certified:
                probe=xb[:min(64,len(xb))]
                net.eval()
                with torch.no_grad(): vold=hidden_energy(net,probe).detach()
                old={k:v.detach().clone() for k,v in net.state_dict().items()}
                # Adam already has a proposal only after step; try decreasing step sizes.
                accepted_this=False; base_lr=opt.param_groups[0]['lr']
                for attempt in range(8):
                    opt.param_groups[0]['lr']=base_lr*(0.5**attempt)
                    opt.step(); opt.zero_grad(set_to_none=True)
                    with torch.no_grad(): vnew=hidden_energy(net,probe).detach()
                    # d is an observed input perturbation scale; certificate is ISS form.
                    d=(probe[:,3:]-probe[:,:-3]).pow(2).mean().sqrt().detach()
                    cert=vnew-vold+lam*vold-c*d*d
                    checks+=1; violations += int(float(cert)>0); cert_margins.append(float(cert))
                    if float(cert)<=0:
                        accepted+=1; accepted_this=True; break
                    net.load_state_dict(old)
                opt.param_groups[0]['lr']=base_lr
                if not accepted_this:
                    rejected+=1
                    # safe fallback: retain the previous parameters, i.e. zero control update.
                    opt.zero_grad(set_to_none=True)
            else:
                opt.step(); opt.zero_grad(set_to_none=True)
            net.train()
    net.eval()
    with torch.no_grad(): metric=float(lossf(net(xte),yte))
    if collect:
        return metric, {'accepted':accepted,'rejected':rejected,'checks':checks,'violations':violations,
                        'violation_rate':violations/max(1,checks),'mean_certificate':float(np.mean(cert_margins)) if cert_margins else 0.0}
    return metric

def baseline_fn(cfg): return lambda seed: train(seed,cfg,False)
def idea_fn(cfg): return lambda seed: train(seed,cfg,True)

def main():
    # Baseline sweep on four seeds, then full paired evaluation; idea 3-config sweep uses same union.
    base=sweep_baseline(baseline_fn,GRID,seeds=SWEEP)
    idea_sweep=[]
    for cfg in IDEA_GRID:
        r=evaluate(idea_fn(cfg),seeds=SWEEP); idea_sweep.append({'cfg':cfg,'mean':r['mean']})
    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']
    idea=evaluate(idea_fn(best),seeds=SEEDS)
    rep=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,
      {'prediction':'Lyapunov certificate filters sampled updates; rejected proposals should have positive certificate and accepted proposals nonpositive.',
       'trained_model_measurements': [train(s,best,True,True)[1] for s in SEEDS],
       'confirmed': False})
    rep['idea']['sweep']=idea_sweep; rep['protocol_notes']='Baseline and idea share rnn_small, data, epochs, batch, Adam, and lr/M grid; only certificate rejection differs.'
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()