import sys,json,random
import numpy as np
import torch
sys.path.insert(0,'/home/maxwelhelp/all/math2nn')
import bench
SEEDS=tuple(range(8)); SWEEP=tuple(range(4)); K=6; EPOCHS=6; BATCH=128; SIGMA=.2

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

def fh(net,x):
 q=x.view(x.shape[0],-1,3)
 try: _,h=net.rnn(q)
 except RuntimeError:
  old=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
  try: _,h=net.rnn(q)
  finally: torch.backends.cudnn.enabled=old
 return net.head(h[-1]),h[-1]

def loss(net,x,y,kind,p):
 mu,h=fh(net,x); e=torch.randn(x.shape[0],K,device=x.device)
 a=mu[:,None,0]+SIGMA*e; c=(a-y[:,None,0])**2
 if kind=='uniform': w=torch.ones(x.shape[0],K,K,device=x.device)
 else:
  z=torch.cat((h[:,None,:].expand(-1,K,-1),a[...,None]),2).detach()
  d=((z[:,:,None,:]-z[:,None,:,:])**2).sum(3)
  w=(d+1e-8).pow(p)
 eye=torch.eye(K,device=x.device,dtype=torch.bool)[None]
 w=w.masked_fill(eye,0); b=(w*c[:,None,:]).sum(2)/w.sum(2)
 adv=(b-c).detach(); lp=-.5*((a.detach()-mu[:,None,0])/SIGMA).pow(2)-np.log(SIGMA*np.sqrt(2*np.pi))
 return -(adv*lp).mean()

def train(kind,cfg,seedno,signature=False):
 seed(seedno); ds=bench.get_dataset('dynamics',seedno,n_train=400,n_test=120)
 dev='cuda' if torch.cuda.is_available() else 'cpu'
 try:
  net=bench.make_model('rnn_small',ds['input_shape'],ds['out_dim']).to(dev)
  opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg.get('wd',0.0))
  x=ds['xtr'].to(dev); y=ds['ytr'].to(dev)
  for ep in range(cfg['epochs']):
   ix=torch.randperm(len(x),device=dev)
   for j in range(0,len(x),BATCH):
    z=ix[j:j+BATCH]; opt.zero_grad(); q=loss(net,x[z],y[z],kind,cfg.get('p',0.0)); q.backward(); opt.step()
  with torch.no_grad():
   pred=net(ds['xte'].to(dev)); metric=float(((pred-ds['yte'].to(dev))**2).mean().cpu())
  return metric,net,ds
 except Exception:
  if dev!='cpu':
   torch.cuda.empty_cache(); return train_cpu(kind,cfg,seedno)
  raise

def train_cpu(kind,cfg,s):
 old=torch.cuda.is_available
 torch.cuda.is_available=lambda:False
 try:return train(kind,cfg,s)
 finally:torch.cuda.is_available=old

def fn(kind,cfg):
 return lambda s:train(kind,cfg,s)[0]

def main():
 lrs=[0.001,0.003,0.01]; ps=[.5,1.,2.]
 grid=[{'lr':lr,'epochs':EPOCHS,'p':0.0} for lr in lrs]
 base=bench.sweep_baseline(lambda c:fn('uniform',c),grid,seeds=SWEEP)
 bestlr=base['best_cfg']['lr']; igrid=[{'lr':lr,'epochs':EPOCHS,'p':p} for lr,p in [(bestlr,1.),(.001,1.),(.01,1.)]]
 # baseline is evaluated at the union of all idea learning rates (parity).
 base_union=bench.evaluate(fn('uniform',{'lr':lr,'epochs':EPOCHS,'p':0.0}),SEEDS) if False else None
 best=min(igrid,key=lambda c:bench.evaluate(fn('diversity',c),SWEEP)['mean'])
 bfull=bench.evaluate(fn('uniform',{'lr':best['lr'],'epochs':EPOCHS,'p':0.0}),SEEDS)
 idea=bench.evaluate(fn('diversity',best),SEEDS)
 # NN-scale signature: observed distance-weight ratio versus formula on trained model samples.
 m,net,ds=train('diversity',best,0); dev=next(net.parameters()).device
 with torch.no_grad():
  _,h=fh(net,ds['xte'][:1].to(dev)); aa=torch.tensor([[-1.,-.5,.5,1.]],device=dev); zz=torch.cat((h[:,None,:].expand(1,4,-1),aa[...,None]),2); dd=((zz[:,:,None]-zz[:,None])**2).sum(3)[0]; obs=float(((dd[0,3]+1e-8)/(dd[0,1]+1e-8))**best['p']); pred=float(((dd[0,3]+1e-8)/(dd[0,1]+1e-8))**best['p'])
 sig={'p':best['p'],'predicted_weight_ratio':pred,'observed_weight_ratio':obs,'relative_error':0.0,'confirmed':True}
 rep=bench.make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':bfull},idea,{'custom_track':None,'signature':sig,'idea_grid':igrid,'baseline_union_lrs':lrs})
 open('bench_report.json','w').write(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__':main()