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,train_model,evaluate,sweep_baseline,make_report,count_params
SEEDS=tuple(range(8)); SWEEP_SEEDS=tuple(range(4)); LRS=[1e-3,3e-3,1e-2]; EPOCHS=20

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 op(n=32):
 L=np.zeros((n,n));
 for i in range(n-1): L[i,i+1]=L[i+1,i]=1
 L=np.diag(L.sum(1))-L; rng=np.random.RandomState(1461); best=None
 for _ in range(30):
  q=rng.normal(size=n); lam,phi=np.linalg.eigh(L+np.diag(q)); S=phi**2
  sums=np.array([lam[i]+lam[j] for i in range(n) for j in range(i,n)])
  gap=np.min(np.abs(sums[:,None]-sums[None,:]+np.eye(len(sums))*1e9)); sm=np.linalg.svd(S,compute_uv=False)[-1]
  if best is None or gap*sm>best[0]: best=(gap*sm,lam,phi,gap,sm)
 return best[1].astype('float32'),best[2].astype('float32'),float(best[3]),float(best[4])

class Base(nn.Module):
 def __init__(self,h=64):
  super().__init__(); self.r=nn.GRU(3,h,batch_first=True); self.head=nn.Linear(h,1)
 def forward(self,x): return self.head(self.r(x.view(x.shape[0],-1,3))[1][-1])

class Spectral(nn.Module):
 def __init__(self,h=64,K=6):
  super().__init__(); self.K=K; self.inp=nn.Linear(3,h); self.h=h
  lam,phi,gap,sm=op(h); self.register_buffer('lam',torch.tensor(lam)); self.register_buffer('phi',torch.tensor(phi));
  self.head=nn.Sequential(nn.Linear(K*h,64),nn.Tanh(),nn.Linear(64,1)); self.gap=gap; self.smin=sm
 def encode(self,z):
  c=z if torch.is_complex(z) else torch.complex(z,torch.zeros_like(z)); out=[]
  for k in range(self.K):
   out.append(torch.sqrt(c.real*c.real+c.imag*c.imag+1e-8)); c=torch.matmul(torch.exp(-1j*self.lam) * torch.matmul(c, self.phi.to(c.dtype)), self.phi.to(c.dtype).T)
  return torch.cat(out,1)
 def forward(self,x):
  z=self.inp(x.view(x.shape[0],-1,3).mean(1)); c=torch.complex(z,torch.zeros_like(z)); out=[]
  for k in range(self.K):
   out.append(torch.sqrt(c.real*c.real+c.imag*c.imag+1e-8)); c=torch.matmul(torch.exp(-1j*self.lam) * torch.matmul(c, self.phi.to(c.dtype)), self.phi.to(c.dtype).T)
  return self.head(torch.cat(out,1))

def run(model,ds,lr,s):
 seed(s); model=model(); _,metric,_=train_model(model,ds,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None); return metric

def main():
 def ds(s): return get_dataset('dynamics',s,n_train=400,n_test=200)
 def mk(c): return lambda s: run(Base,ds(s),c['lr'],s)
 base=sweep_baseline(mk,[{'lr':x} for x in LRS],seeds=SWEEP_SEEDS)
 # Full evaluation for every union lr, with the selected best retained as the idea comparison.
 base_full={}
 for lr in LRS: base_full[lr]=evaluate(mk({'lr':lr}),SEEDS)
 best_lr=base['best_cfg']['lr']; idea_cfgs=[best_lr]+[x for x in LRS if x!=best_lr][:2]
 ir=[]
 for lr in idea_cfgs: ir.append((lr,evaluate(lambda s:run(lambda:Spectral(K=6),ds(s),lr,s),SEEDS)))
 idea_lr,idea=min(ir,key=lambda z:z[1]['mean'])
 d=ds(0); m=Spectral(K=6); seed(0); _,_,_=train_model(m,d,epochs=EPOCHS,lr=idea_lr,batch=128,log=lambda *_:None)
 m.eval(); dev=next(m.parameters()).device; x=d['xte'][:32].to(dev)
 with torch.no_grad():
  z=m.inp(x.view(x.shape[0],-1,3).mean(1)); e1=m.encode(z); e2=m.encode(torch.complex(z,torch.zeros_like(z))*torch.exp(1j*torch.tensor(1.234)))
 sig={'prediction':'global phase is unobservable by magnitude trajectory','predicted_max_abs_difference':0.0,'observed_mean_abs_difference':float((e1-e2).abs().mean()),'operator_pair_sum_gap':m.gap,'operator_sigma_min_S':m.smin,'confirmed':bool(float((e1-e2).abs().mean())<1e-5)}
 report=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base_full[best_lr]},idea,{'mechanism_signature':sig})
 report['idea_sweep']=[{'cfg':{'lr':lr},'result':r} for lr,r in ir]; report['parameter_counts']={'baseline':count_params(Base()),'idea':count_params(Spectral())}
 with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
 print(json.dumps(report,indent=2))
if __name__=='__main__': main()