import sys, json, random
from pathlib import Path
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, train_model, evaluate, sweep_baseline, make_report
TRACK='dynamics'; MODEL='rnn_small'; EPOCHS=4; NTR=240; NTE=240; BATCH=64
LRS=[1e-3,3e-3,1e-2]; ALPHAS=[1e-4,3e-4,1e-3]

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

def fisher_loss(net,x,nsamp=2,eps=1e-3):
    # Fisher proxy for the learned recurrent rollout: output hidden trajectory
    # sensitivity to the initial observed state (theta, omega).
    xx=x[:nsamp].detach().clone().requires_grad_(True)
    seq=xx.view(xx.shape[0],-1,3); out,_=net.rnn(seq)
    vals=[]
    for i in range(nsamp):
        rows=[]
        # one Jacobian row per observed rollout time, preserving the chain rule
        for t in range(out.shape[1]):
            g=torch.autograd.grad(out[i,t,0],xx,retain_graph=True,create_graph=True)[0][i,:2]
            rows.append(g)
        J=torch.stack(rows); I=J.T@J+eps*torch.eye(2,device=x.device)
        ev=torch.linalg.eigvalsh(I)
        vals.append(-torch.logdet(I)+1e-3*ev[-1]/(ev[0]+eps))
    return torch.stack(vals).mean()

def idea_train(seed,lr,alpha,return_net=False):
    seedall(seed); ds=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE)
    net=make_model(MODEL,ds['input_shape'],ds['out_dim'])
    try:
        device='cuda' if torch.cuda.is_available() else 'cpu'; net.to(device)
        x,y=ds['xtr'].to(device),ds['ytr'].to(device); opt=torch.optim.Adam(net.parameters(),lr=lr); mse=nn.MSELoss(); gen=torch.Generator().manual_seed(seed)
        for _ in range(EPOCHS):
            net.train()
            for ix in torch.randperm(len(x),generator=gen).split(BATCH):
                xb,yb=x[ix],y[ix]; opt.zero_grad(set_to_none=True)
                pred=net(xb); loss=mse(pred,yb)+alpha*fisher_loss(net,xb); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5); opt.step()
        net.eval()
        with torch.no_grad(): metric=float(mse(net(ds['xte'].to(device)),ds['yte'].to(device)).cpu())
        return (metric,net,ds,device) if return_net else metric
    except Exception:
        # CPU fallback for shared/limited CUDA environments.
        seedall(seed); ds=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE); net=make_model(MODEL,ds['input_shape'],ds['out_dim']); opt=torch.optim.Adam(net.parameters(),lr=lr); mse=nn.MSELoss(); gen=torch.Generator().manual_seed(seed)
        for _ in range(EPOCHS):
            for ix in torch.randperm(len(ds['xtr']),generator=gen).split(BATCH):
                xb,yb=ds['xtr'][ix],ds['ytr'][ix]; opt.zero_grad(); loss=mse(net(xb),yb)+alpha*fisher_loss(net,xb); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5); opt.step()
        with torch.no_grad(): metric=float(mse(net(ds['xte']),ds['yte']))
        return (metric,net,ds,'cpu') if return_net else metric

def signature():
    ans={}
    for name,a in [('baseline',0.0),('idea',3e-4)]:
        metric,net,ds,device=idea_train(0,3e-3,a,True); x=ds['xte'][:4].to(device).requires_grad_(True); out,_=net.rnn(x.view(4,-1,3)); evs=[]
        for i in range(4):
            J=torch.stack([torch.autograd.grad(out[i,t,0],x,retain_graph=True)[0][i,:2] for t in range(out.shape[1])]); evs.append(torch.linalg.eigvalsh(J.T@J+1e-3*torch.eye(2,device=device)).detach().cpu().numpy())
        evs=np.asarray(evs); ans[name]={'test_mse_observed':metric,'fisher_lambda_min_predicted':float(evs[:,0].mean()),'fisher_lambda_max_predicted':float(evs[:,1].mean()),'condition_predicted':float(np.mean(evs[:,1]/evs[:,0]))}
    ans['confirmed']=ans['idea']['fisher_lambda_min_predicted']>ans['baseline']['fisher_lambda_min_predicted'] and ans['idea']['condition_predicted']<ans['baseline']['condition_predicted']; return ans

def main():
    # Baseline gets the full union of all LRs used by the idea side.
    base=sweep_baseline(lambda c: lambda s: float(train_model(make_model(MODEL,get_dataset(TRACK,s,NTR,NTE)['input_shape'],1),get_dataset(TRACK,s,NTR,NTE),epochs=EPOCHS,lr=c['lr'],batch=BATCH,log=lambda *_:None)[1]),[{'lr':v} for v in LRS])
    idea_runs=[]
    for a in ALPHAS:
        r=evaluate(lambda s,a=a: idea_train(s,base['best_cfg']['lr'],a)); idea_runs.append((r,a))
    idea,a=min(idea_runs,key=lambda z:z[0]['mean'])
    rep=make_report(TRACK,MODEL,base,idea,{'type':'trained_model_fisher_signature','values':signature(),'confirmed_definition':'true only when trained idea increases mean predicted lambda_min and lowers condition number'})
    rep['idea_sweep']=[{'alpha':a,'full':r} for r,a in idea_runs]; rep['protocol_notes']='Dynamics is the structural match: controlled pendulum multi-step rollout. Same rnn_small, data, epochs, batch, and baseline-selected LR; only Fisher loss differs.'
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()