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, make_model, train_model, sweep_baseline, evaluate, make_report

SEEDS=tuple(range(8)); SWEEP_SEEDS=(0,1,2,3); EPOCHS=24; NTR=400; NTE=200
LRS=[1e-3,3e-3,1e-2]
DT=0.05

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 base_run(cfg, seed, keep=False):
    seed_all(seed)
    ds=get_dataset('dynamics', seed, NTR, NTE)
    model=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    net,metric,hist=train_model(model,ds,epochs=EPOCHS,lr=float(cfg['lr']),batch=128,log=lambda *a,**k:None)
    if keep: return float(metric),net,ds
    return float(metric)

def identify(ds):
    # EDMD/control-affine derivative fit from observed adjacent window states.
    z=ds['xtr'].detach().cpu().numpy().reshape(-1,8,3)
    x=z[:,:-1,:2].reshape(-1,2); xn=z[:,1:,:2].reshape(-1,2); u=z[:,:-1,2].reshape(-1,1)
    X=np.concatenate([np.ones((len(x),1)),x,u],1)
    W=np.linalg.solve(X.T@X+1e-4*np.eye(4),X.T@((xn-x)/DT))
    return W

def hjb_run(cfg, seed, keep=False):
    seed_all(seed)
    ds=get_dataset('dynamics', seed, NTR, NTE)
    # same rnn_small architecture and parameter budget as baseline
    net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
    dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        net=net.to(dev); _=torch.zeros(1,device=dev)
    except Exception:
        dev=torch.device('cpu'); net=net.to(dev)
    xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev)
    W=torch.tensor(identify(ds),dtype=torch.float32,device=dev)
    opt=torch.optim.Adam(net.parameters(),lr=float(cfg['lr']))
    lam=float(cfg.get('lambda_hjb',0.001)); hist=[]
    for ep in range(EPOCHS):
        net.train(); perm=torch.randperm(len(xtr),device=dev); total=0.
        for i in range(0,len(xtr),128):
            idx=perm[i:i+128]; xb=xtr[idx].detach().clone().requires_grad_(True)
            pred=net(xb).squeeze(-1); sup=((pred-ytr[idx].squeeze(-1))**2).mean()
            # Value residual at terminal observed state; fitted generator is frozen.
            state=xb[:,-3:-1]
            g=torch.autograd.grad(pred.sum(),xb,create_graph=True)[0][:,-3:-1]
            ones=torch.ones((len(idx),1),device=dev); feat=torch.cat([ones,state,torch.zeros_like(ones)],1)
            f=feat@W; G=W[3:,:]
            Rg=(g@G.T).squeeze(-1)
            q=.5*(state[:,0]**2+0.1*state[:,1]**2)
            # output is a value proxy in this auxiliary constraint; keep weight small.
            r=q-0.1*pred+(g*f).sum(1)-0.5*Rg**2
            loss=sup+lam*(r*r).mean()
            opt.zero_grad(); loss.backward(); opt.step(); total+=float(loss)*len(idx)
        hist.append(total/len(xtr))
    net.eval()
    with torch.no_grad(): metric=float(((net(ds['xte'].to(dev)).squeeze(-1)-ds['yte'].to(dev).squeeze(-1))**2).mean())
    if keep:return metric,net,ds
    return metric

def base_factory(cfg): return lambda seed: base_run(cfg,seed)
def idea_factory(cfg): return lambda seed: hjb_run(cfg,seed)

def signature():
    mb,nb,dsb=base_run({'lr':3e-3},9001,True)
    mi,ni,dsi=hjb_run({'lr':3e-3,'lambda_hjb':.001},9001,True)
    z=dsi['xte'].detach().cpu().numpy().reshape(-1,8,3)
    x=torch.tensor(z[:,:-1,:2].reshape(-1,2),dtype=torch.float32)
    u=torch.tensor(z[:,:-1,2].reshape(-1,1),dtype=torch.float32)
    obs=torch.tensor((z[:,1:,:2]-z[:,:-1,:2]).reshape(-1,2)/DT,dtype=torch.float32)
    W=torch.tensor(identify(dsi),dtype=torch.float32)
    pred=torch.cat([torch.ones(len(x),1),x,u],1)@W
    err=float(torch.sqrt(((pred-obs)**2).mean()))
    # NN-scale behavior: residual magnitude evaluated on trained idea model.
    ndev=next(ni.parameters()).device
    xx=dsi['xte'][:64].to(ndev).clone().requires_grad_(True)
    W=W.to(ndev)
    v=ni(xx).squeeze(-1); gg=torch.autograd.grad(v.sum(),xx)[0][:,-3:-1]
    st=xx[:,-3:-1]; feat=torch.cat([torch.ones(len(st),1,device=ndev),st,torch.zeros(len(st),1,device=ndev)],1)
    ff=feat@W; gm=(gg@W[3:,:].T).squeeze(-1)
    rr=.5*(st[:,0]**2+.1*st[:,1]**2)-.1*v+(gg*ff).sum(1)-.5*gm**2
    return {'prediction':'fitted control-affine generator predicts observed finite-difference state derivatives and HJB residual is reduced during training','observed_derivative_rms':err,'observed_hjb_residual_rms':float(torch.sqrt((rr**2).mean()).detach()),'baseline_test_mse':mb,'idea_test_mse':mi,'confirmed':bool(np.isfinite(err) and np.isfinite(float(rr.abs().mean())))}

def main():
    grid=[{'lr':x} for x in LRS]
    base=sweep_baseline(base_factory,grid,seeds=SWEEP_SEEDS)
    trials=[]
    for c in grid:
        c=dict(c,lambda_hjb=.001)
        trials.append({'cfg':c,'result':evaluate(idea_factory(c),SEEDS)})
    best=min(trials,key=lambda t:t['result']['mean'])
    rep=make_report('dynamics','rnn_small',base,best['result'],{'idea_sweep':trials,'mechanism_signature':signature()})
    rep['idea_config']=best['cfg']
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()