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

SEEDS = tuple(range(8)); EPOCHS = 5; BATCH = 128; H = 8
LRS = (0.0015, 0.003, 0.006); LAMBDAS = (0.01, 0.03, 0.10)
CACHE = {}

class DynamicsGRU(nn.Module):
    def __init__(self, hidden=16):
        super().__init__(); self.rnn=nn.GRU(3,hidden,batch_first=True); self.head=nn.Linear(hidden,1)
    def forward(self,x,h=None,return_seq=False):
        out,hn=self.rnn(x.view(x.shape[0],-1,3),h); pred=self.head(hn[-1])
        return (pred,out,hn) if return_seq else pred

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

def train_one(seed,lr,lam):
    key=(int(seed),float(lr),float(lam))
    if key in CACHE: return CACHE[key]
    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
    requested='cuda' if torch.cuda.is_available() else 'cpu'
    for dev in ([requested,'cpu'] if requested=='cuda' else ['cpu']):
        try:
            model=DynamicsGRU().to(dev); xtr=ds['xtr'].to(dev); ytr=ds['ytr'].to(dev).view(-1,1)
            opt=torch.optim.Adam(model.parameters(),lr=lr); mse=nn.MSELoss()
            gen=torch.Generator(device=dev); gen.manual_seed(seed+10000)
            for _ in range(EPOCHS):
                order=torch.randperm(len(xtr),generator=gen,device=dev)
                for st in range(0,len(xtr),BATCH):
                    ix=order[st:st+BATCH]; xb=xtr[ix]; task=mse(model(xb),ytr[ix])
                    if lam:
                        b=len(ix); h0=torch.zeros(1,b,16,device=dev)
                        eps=torch.randn(b,16,generator=gen,device=dev)*.02
                        _,a,_=model(xb,h0,True); _,bseq,_=model(xb,h0+eps.unsqueeze(0),True)
                        roll=(a-bseq).pow(2).sum(-1).sqrt().amax(1).div(eps.norm(dim=1)+1e-8).mean()
                        loss=task+lam*roll
                    else: loss=task
                    opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
            model.eval()
            with torch.no_grad():
                xt=ds['xte'].to(dev); yt=ds['yte'].to(dev).view(-1,1)
                metric=float(((model(xt)-yt)**2).mean())
                b=len(xt); h0=torch.zeros(1,b,16,device=dev)
                eps=torch.randn(b,16,generator=gen,device=dev)*.02
                _,a,_=model(xt,h0,True); _,bb,_=model(xt,h0+eps.unsqueeze(0),True)
                gain=float(((a-bb).pow(2).sum(-1).sqrt().amax(1)/(eps.norm(dim=1)+1e-8)).median())
            CACHE[key]=(metric,gain); return CACHE[key]
        except RuntimeError:
            if dev=='cuda': continue
            raise
    raise RuntimeError('training failed')

def fn(cfg): return lambda s: train_one(s,cfg['lr'],cfg.get('lambda',0.0))[0]

def main():
    # Baseline covers every LR used by the idea-side shared architecture.
    base=sweep_baseline(fn,[{'lr':lr,'lambda':0.0} for lr in LRS])
    best_lr=float(base['best_cfg']['lr'])
    idea_runs=[]
    for lam in LAMBDAS:
        cfg={'lr':best_lr,'lambda':lam}; idea_runs.append((evaluate(fn(cfg),SEEDS),cfg))
    idea,idea_cfg=min(idea_runs,key=lambda z:z[0]['mean'])
    base_cfg=base['best_cfg']
    sb=[train_one(s,base_cfg['lr'],0.0)[1] for s in SEEDS]
    si=[train_one(s,idea_cfg['lr'],idea_cfg['lambda'])[1] for s in SEEDS]
    red=1-float(np.mean(si))/float(np.mean(sb))
    sig={'quantity':'median test-set G_H=max hidden separation / initial perturbation norm','H':H,
      'predicted':'rollout penalty reduces finite-horizon gain','baseline_mean_GH':float(np.mean(sb)),
      'idea_mean_GH':float(np.mean(si)),'observed_relative_reduction':float(red),
      'predicted_vs_observed':{'predicted_direction':'decrease','observed_direction':'decrease' if red>0 else 'increase'},
      'confirmed':bool(red>0.10),'paired_seed_GH_baseline':sb,'paired_seed_GH_idea':si,
      'idea_config':idea_cfg,'track_choice_justification':'dynamics is the built-in control/stability track with multi-step pendulum rollouts.'}
    rep=make_report('dynamics','rnn_small',base,idea,extra=sig)
    rep['idea_sweep']=[{'cfg':c,'mean':r['mean'],'per_seed':r['per_seed']} for r,c in idea_runs]
    rep['protocol_notes']={'epochs':EPOCHS,'n_train':400,'n_test':200,'baseline_lr_union':list(LRS),'idea_lambdas':list(LAMBDAS),'cache_used':True}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()