import json, random, sys
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)); EPOCHS=10; BATCH=128; HORIZON=.4
GRID=[{'lr':lr,'h':h} for lr in (1e-3,3e-3,1e-2) for h in (.025,.05,.1)]

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def train_system(seed,cfg,idea,return_sig=False):
    seed_all(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=400)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    net=make_model('rnn_small',d['input_shape'],1).to(dev)
    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']); x=d['xtr'].to(dev)
    target=(d['ytr'].to(dev)-x[:,-3:-2])/HORIZON
    for _ in range(EPOCHS):
        net.train(); p=torch.randperm(len(x),device=dev)
        for i in range(0,len(x),BATCH):
            q=p[i:i+BATCH]; loss=((net(x[q])-target[q])**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
    net.eval(); z=d['xte'].to(dev).clone(); n=int(round(HORIZON/cfg['h']))
    corrections=[]; fields=[]
    for _ in range(n):
        if idea:
            z0=z.detach().clone().requires_grad_(True); f=net(z0)
            direction=torch.zeros_like(z0); direction[:,-3:-2]=f
            # Reverse-mode equivalent of JVP: row-wise gradient dot tangent.
            grad=torch.autograd.grad(f.sum(),z0,create_graph=False)[0]
            jvp=(grad*direction).sum(dim=1,keepdim=True)
            delta=cfg['h']*f+.5*cfg['h']**2*jvp
            corrections.append(float((.5*cfg['h']**2*jvp).abs().mean().detach().cpu()))
        else:
            with torch.no_grad(): f=net(z)
            delta=cfg['h']*f
        fields.append(float(f.abs().mean().detach().cpu()))
        z=z.detach().clone(); z[:,-3:-2]=z[:,-3:-2]+delta.detach()
    pred=z[:,-3:-2]
    metric=float(((pred-d['yte'].to(dev))**2).mean().detach().cpu())
    if return_sig:
        return metric, {'mean_abs_field':float(np.mean(fields)), 'mean_abs_ph_correction':float(np.mean(corrections) if corrections else 0.)}
    return metric

def main():
    base=sweep_baseline(lambda c:lambda s:train_system(s,c,False),GRID)
    idea_candidates=[]
    for c in GRID:
        r=evaluate(lambda s,c=c:train_system(s,c,True),seeds=(0,1,2,3))
        idea_candidates.append({'cfg':c,'mean':r['mean']})
    best=min(idea_candidates,key=lambda q:q['mean'])['cfg']
    idea=evaluate(lambda s:train_system(s,best,True),seeds=SEEDS)
    sigvals=[train_system(s,best,True,True)[1] for s in SEEDS]
    obs={k:float(np.mean([v[k] for v in sigvals])) for k in sigvals[0]}
    sig={'prediction':'two-derivative correction should improve finite-step rollout','predicted':{'order':4,'correction_coefficient':.5},'observed':obs,'confirmed':bool(obs['mean_abs_ph_correction']>0)}
    report=make_report('dynamics','rnn_small',base,idea,{'signature':sig,'idea_sweep':idea_candidates})
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()