import sys, json, math, 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, sweep_baseline, evaluate, make_report

SEED0=1729
DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
N=16; TAU_R=1.0; TAU_D=0.10
MODES=(1,2,3,-1,-2,-3)

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 root_np(what,a=1.,tr=1.,td=TAU_D,it=20):
    z=complex((-1+a*what)/tr)
    for _ in range(it):
        e=np.exp(-z*td); f=tr*z+1-a*what*e; df=tr+a*what*td*e
        if abs(df)<1e-12: break
        z-=f/df
    return z

def sanity():
    # FFT formula and characteristic-root prediction against a direct Euler rollout.
    rng=np.random.default_rng(3); w=rng.normal(size=N)*.04
    hats=np.fft.fft(w); k=2; explicit=sum(w[j]*np.exp(-2j*np.pi*k*j/N) for j in range(N))
    what=hats[k]; lam=root_np(what)
    dt=.001; delay=round(TAU_D/dt); steps=3500
    x=np.zeros(steps+delay+1,dtype=complex); x[:delay+1]=1e-3*np.exp(lam*np.arange(-delay,1)*dt)
    for i in range(delay,delay+steps): x[i+1]=x[i]+dt*(-x[i]+what*x[i-delay])
    t=np.arange(steps)*dt; sel=t>1.0
    measured=np.polyfit(t[sel],np.log(np.abs(x[delay:delay+steps][sel])+1e-30),1)[0]
    return {'fft_error':float(abs(what-explicit)), 'predicted_growth':float(lam.real),
            'observed_growth':float(measured), 'abs_error':float(abs(lam.real-measured)),
            'passed':bool(abs(what-explicit)<1e-10 and abs(lam.real-measured)<0.02)}

class RingRNN(nn.Module):
    def __init__(self, hidden=N):
        super().__init__(); self.hidden=hidden
        self.inp=nn.Linear(3,hidden); self.w=nn.Parameter(torch.randn(hidden)*0.05)
        self.out=nn.Linear(hidden,1)
    def recurrent(self,h):
        # circulant convolution; FFT convention matches the proposed w_hat formula.
        return torch.fft.ifft(torch.fft.fft(self.w)*torch.fft.fft(h,dim=-1),dim=-1).real
    def forward(self,x,return_aux=False):
        q=x.view(x.shape[0],-1,3); h=torch.zeros(x.shape[0],self.hidden,device=x.device)
        slopes=[]; preacts=[]
        for t in range(q.shape[1]):
            p=self.inp(q[:,t])+self.recurrent(h)
            slopes.append(1-torch.tanh(p).pow(2)); preacts.append(p)
            h=torch.tanh(p)
        y=self.out(h)
        if return_aux: return y, torch.stack(slopes).mean(), torch.stack(preacts)
        return y

def spectral_penalty(net,a):
    hats=torch.fft.fft(net.w)
    vals=[]
    roots=[]
    for k in MODES:
        z=(-1+a*hats[k])/TAU_R
        for _ in range(8):
            e=torch.exp(-z*TAU_D); f=TAU_R*z+1-a*hats[k]*e
            z=z-f/(TAU_R+a*hats[k]*TAU_D*e)
        roots.append(z); vals.append(torch.nn.functional.softplus(z.real+0.02)**2)
    return torch.stack(vals).mean(), torch.stack(roots)

def train_one(seed, lr, wd, strength):
    seed_all(seed)
    d=get_dataset('dynamics',seed,n_train=1000,n_test=300)
    net=RingRNN().to(DEVICE)
    opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd)
    xtr,ytr=d['xtr'].to(DEVICE),d['ytr'].to(DEVICE)
    lossf=nn.MSELoss(); batch=128; last_roots=None
    for _ in range(12):
        net.train(); perm=torch.randperm(len(xtr),device=DEVICE)
        for i in range(0,len(xtr),batch):
            ix=perm[i:i+batch]; pred,a,_=net(xtr[ix],True); loss=lossf(pred,ytr[ix].view(-1,1))
            if strength>0:
                sp,last_roots=spectral_penalty(net,a.detach()); loss=loss+strength*sp
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
    net.eval()
    with torch.no_grad(): metric=float(lossf(net(d['xte'].to(DEVICE)),d['yte'].to(DEVICE).view(-1,1)).cpu())
    # Signature uses the trained model: predicted characteristic growth versus
    # observed infinitesimal autonomous perturbation growth of its recurrence.
    with torch.no_grad():
        a=float(net.inp.weight.new_tensor(0.0))
        q=d['xte'][:64].to(DEVICE).view(64,-1,3); h=torch.zeros(64,N,device=DEVICE)
        ss=[]
        for t in range(q.shape[1]):
            p=net.inp(q[:,t])+net.recurrent(h); ss.append((1-torch.tanh(p).pow(2)).mean()); h=torch.tanh(p)
        a=float(torch.stack(ss).mean().cpu()); hats=torch.fft.fft(net.w).cpu().numpy()
        pred=max(root_np(h,a).real for h in hats[[k%N for k in MODES]])
        h0=h[:1].clone(); eps=1e-4; hp=h0+eps*torch.randn_like(h0)
        vals=[]
        for _ in range(30):
            hp=torch.tanh(net.recurrent(hp)); vals.append(float(torch.linalg.vector_norm(hp-h0).cpu())+1e-30)
        obs=float(np.polyfit(np.arange(30),np.log(vals),1)[0])
    return metric, {'predicted_max_growth':float(pred),'observed_perturbation_growth':obs,'local_slope':a,
                    'growth_abs_error':abs(float(pred)-obs)}

def main():
    sanity_result=sanity(); print('SANITY',json.dumps(sanity_result))
    grid=[{'lr':lr,'wd':wd} for lr in (1e-3,3e-3,6e-3) for wd in (0.,1e-4)]
    def base_fn(cfg): return lambda s: train_one(s,cfg['lr'],cfg['wd'],0.)[0]
    base=sweep_baseline(base_fn,grid,seeds=(0,1,2,3))
    best=base['best_cfg']; strengths=(0.01,0.03,0.10)
    idea_cfgs=[{'lr':best['lr'],'wd':best['wd'],'strength':v} for v in strengths]
    idea_trials=[]
    for cfg in idea_cfgs:
        r=evaluate(lambda s: train_one(s,cfg['lr'],cfg['wd'],cfg['strength'])[0])
        idea_trials.append((cfg,r))
    cfg,idea=min(idea_trials,key=lambda z:z[1]['mean'])
    # Re-run best idea at all eight seeds already done above; collect signatures separately.
    sig=[train_one(s,cfg['lr'],cfg['wd'],cfg['strength'])[1] for s in range(8)]
    sigmean={k:float(np.mean([z[k] for z in sig])) for k in sig[0]}
    sigmean['confirmed']=bool(sigmean['growth_abs_error']<0.15)
    report=make_report('dynamics','ring_rnn',base,idea,{'predicted_vs_observed_growth':sigmean,
        'stage1_sanity':sanity_result,'idea_sweep':[{'cfg':c,'mean':r['mean']} for c,r in idea_trials]})
    report['parameter_count']=sum(p.numel() for p in RingRNN().parameters())
    report['device']=DEVICE
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()