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

N, M, KAPPA, Q = 32, 4, 0.15, 0
SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
# Shared union: every idea learning rate is also evaluated by baseline.
GRID = [{'lr': 1e-3, 'weight_decay': 0.0},
        {'lr': 3e-3, 'weight_decay': 0.0},
        {'lr': 1e-2, 'weight_decay': 0.0}]


def mu_factors(weights, q=Q):
    w = np.asarray(weights, dtype=float)
    ell = np.arange(1, len(w)+1)
    cq = np.cos(2*np.pi*q*ell/N)
    return np.array([np.sum(w*cq*(1-np.cos(2*np.pi*k*ell/N))) for k in range(N)])

class CyclicRNN(nn.Module):
    def __init__(self, hidden=N):
        super().__init__()
        self.inp = nn.Linear(3, hidden)
        self.weights = nn.Parameter(torch.ones(M))
        self.head = nn.Linear(hidden, 1)
    def transition(self, h, z):
        v = torch.zeros_like(h)
        for j in range(M):
            # forward neighbor i+j+1, matching the paper's convention
            v = v + self.weights[j] * (torch.roll(h, -(j+1), dims=1) - h)
        return torch.tanh(h + KAPPA*v + z)
    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        h = torch.zeros(x.shape[0], N, device=x.device, dtype=x.dtype)
        for t in range(seq.shape[1]):
            h = self.transition(h, self.inp(seq[:, t]))
        return self.head(h)

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

def train_idea(model, ds, epochs, lr, weight_decay=0.0, lam=0.03, eps=0.02):
    errs=[]
    ladder = [('cuda', False), ('cuda', True)] if torch.cuda.is_available() else []
    ladder.append(('cpu', False))
    for device, no_cudnn in ladder:
        try:
            if no_cudnn: torch.backends.cudnn.enabled=False
            net=model.to(device); opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=weight_decay)
            x,y=ds['xtr'].to(device),ds['ytr'].to(device)
            for _ in range(epochs):
                net.train(); perm=torch.randperm(len(x),device=device)
                for i in range(0,len(x),128):
                    ix=perm[i:i+128]; pred=net(x[ix]); loss=((pred-y[ix])**2).mean()
                    mu=torch.stack([torch.sum(net.weights*torch.tensor(
                        np.cos(2*np.pi*Q*np.arange(1,M+1)/N)*(1-np.cos(2*np.pi*k*np.arange(1,M+1)/N)),
                        device=device,dtype=net.weights.dtype)) for k in range(1,N)])
                    barrier=torch.nn.functional.softplus(eps-mu).pow(2).mean()
                    loss=loss+lam*barrier
                    opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),10.0); opt.step()
            net.eval()
            with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
            if no_cudnn: torch.backends.cudnn.enabled=True
            return net, metric
        except RuntimeError as e:
            errs.append(str(e)[:100])
            if no_cudnn: torch.backends.cudnn.enabled=True
    return None, float('nan')

def base_make(cfg):
    def fn(seed):
        seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
        net=CyclicRNN(); _,metric,_=train_model(net,ds,epochs=18,lr=cfg['lr'],batch=128,weight_decay=cfg['weight_decay'],log=lambda *_:None)
        return metric
    return fn

def idea_make(cfg):
    def fn(seed):
        seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
        _,metric=train_idea(CyclicRNN(),ds,epochs=18,lr=cfg['lr'],weight_decay=cfg['weight_decay'])
        return metric
    return fn

def choose_idea():
    tried=[]
    for cfg in GRID:
        r=evaluate(idea_make(cfg),SWEEP_SEEDS); tried.append({'cfg':cfg,'mean':r['mean']})
    best=min(tried,key=lambda x:x['mean'])['cfg']
    return {'best_cfg':best,'sweep':tried,'full':evaluate(idea_make(best),SEEDS)}

def signature(cfg):
    seed_all(0); ds=get_dataset('dynamics',0,n_train=400,n_test=200)
    b=CyclicRNN(); b, bm,_=train_model(b,ds,epochs=18,lr=cfg['lr'],batch=128,log=lambda *_:None)
    seed_all(0); i,im=train_idea(CyclicRNN(),ds,epochs=18,lr=cfg['lr'])
    rows=[]
    for name,net in [('baseline',b),('idea',i)]:
        w=net.weights.detach().cpu().numpy(); mu=mu_factors(w)
        obs=[]
        with torch.no_grad():
            for k in (1,2,3):
                phase=2*np.pi*k*np.arange(N)/N
                h=(1e-4*torch.tensor(np.cos(phase),dtype=torch.float32)[None,:]).to(next(net.parameters()).device)
                amps=[]
                for _ in range(30):
                    z=torch.zeros_like(h); h=net.transition(h,z)
                    amps.append(float(torch.abs(torch.fft.fft(h)[0,k]).cpu()))
                slope=float(np.polyfit(np.arange(len(amps))*1.0,np.log(np.maximum(amps,1e-30)),1)[0])
                pred=float(-KAPPA*mu[k]); obs.append({'mode':k,'predicted':pred,'observed':slope})
        rows.append({'system':name,'weights':w.tolist(),'modes':obs})
    rel=[abs(x['observed']-x['predicted'])/max(abs(x['predicted']),1e-6) for r in rows for x in r['modes']]
    return {'q':Q,'kappa':KAPPA,'rows':rows,'max_relative_error':float(max(rel)),'confirmed':bool(max(rel)<0.20),'task_metric_seed0':{'baseline':bm,'idea':im}}

def main():
    base=sweep_baseline(base_make,GRID,seeds=SWEEP_SEEDS)
    idea=choose_idea()
    # Ensure idea uses baseline-selected lr if its sweep happened to select another: report its best, all are shared.
    rep=make_report('dynamics','cyclic_rnn',base,idea['full'],{'note':'trained-model zero-input Fourier perturbation; q=0','signature':signature(idea['best_cfg'])})
    rep['idea_sweep']=idea['sweep']; rep['protocol']={'paired_seeds':list(SEEDS),'train_samples':400,'test_samples':200,'epochs':18,'architecture':'same CyclicRNN; barrier only differs'}
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()