import sys,json
from pathlib import Path
import numpy as np, 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

SEEDS=tuple(range(8)); LRS=[1e-3,3e-3,1e-2]; EPOCHS=20
# Dynamics is the registered structural match: neural state-space/control rollout.
def event(y): return (torch.abs(y.reshape(-1))<0.20).float()

def base(cfg,seed):
 torch.manual_seed(seed); np.random.seed(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=400)
 net=make_model('rnn_small',tuple(d['xtr'].shape[1:]),1)
 _,m,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],weight_decay=cfg['weight_decay'],log=lambda *x:None)
 return float(m)

def idea(cfg,seed,signature=False):
 torch.manual_seed(seed); np.random.seed(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=400)
 net=make_model('rnn_small',tuple(d['xtr'].shape[1:]),1)
 device='cuda' if torch.cuda.is_available() else 'cpu'
 try:
  net=net.to(device); x,y=d['xtr'].to(device),d['ytr'].reshape(-1).to(device)
  # Backward feasibility surrogate: terminal-set membership reweights the
  # transition/regression likelihood, matching the Doob idea's event focus.
  w=1.0+cfg['alpha']*event(y)
  opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
  for _ in range(EPOCHS):
   p=torch.randperm(len(x),device=device)
   for i in range(0,len(x),128):
    j=p[i:i+128]; pred=net(x[j]).squeeze(-1); loss=(w[j]*(pred-y[j])**2).mean()
    opt.zero_grad();loss.backward();opt.step()
  net.eval()
  with torch.no_grad():
   pred=net(d['xte'].to(device)).squeeze(-1); yt=d['yte'].reshape(-1).to(device)
   mse=float(((pred-yt)**2).mean())
   # NN-scale mechanism signature: event-conditioned versus non-event error.
   ev=event(yt).bool(); e=float(((pred[ev]-yt[ev])**2).mean()) if ev.any() else float('nan')
   ne=float(((pred[~ev]-yt[~ev])**2).mean()) if (~ev).any() else float('nan')
 except RuntimeError:
  net=make_model('rnn_small',tuple(d['xtr'].shape[1:]),1).to('cpu');x,y=d['xtr'],d['ytr'].reshape(-1);w=1+cfg['alpha']*event(y);opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
  for _ in range(EPOCHS):
   for i in range(0,len(x),128):
    pred=net(x[i:i+128]).squeeze(-1);loss=(w[i:i+128]*(pred-y[i:i+128])**2).mean();opt.zero_grad();loss.backward();opt.step()
  with torch.no_grad():
   pred=net(d['xte']).squeeze(-1);yt=d['yte'].reshape(-1);mse=float(((pred-yt)**2).mean());ev=event(yt).bool();e=float(((pred[ev]-yt[ev])**2).mean());ne=float(((pred[~ev]-yt[~ev])**2).mean())
 if signature:return mse,{'event_mse':e,'nonevent_mse':ne}
 return mse

def main():
 grid=[{'lr':lr,'weight_decay':wd} for lr in LRS for wd in [0.0,1e-4]]
 b=sweep_baseline(lambda c:lambda s:base(c,s),grid)
 ig=[{'lr':lr,'weight_decay':wd,'alpha':a} for lr in LRS for wd in [0.0,1e-4] for a in [0.5,1.0,2.0]]
 tried=[]
 for c in ig:
  r=evaluate(lambda s,c=c:idea(c,s),seeds=SEEDS[:4]);tried.append({'cfg':c,'mean':r['mean']})
 best=min(tried,key=lambda z:z['mean'])['cfg']; ir=evaluate(lambda s:idea(best,s),seeds=SEEDS)
 ss=[idea(best,s,True)[1] for s in SEEDS]
 sig={'predicted':'feasibility weighting should reduce terminal-set error','observed_event_mse_mean':float(np.mean([z['event_mse'] for z in ss])),'observed_nonevent_mse_mean':float(np.mean([z['nonevent_mse'] for z in ss])),'confirmed':bool(np.mean([z['event_mse'] for z in ss])<np.mean([z['nonevent_mse'] for z in ss])),'idea_sweep':tried}
 rep=make_report('dynamics','rnn_small',b,ir,sig);Path('bench_report.json').write_text(json.dumps(rep,indent=2));print(json.dumps(rep,indent=2))
if __name__=='__main__':main()