import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED=2540
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

# Toy verification: f_a(x)=a*tanh(x), U=[-alpha,alpha].
# Since |tanh(x)| is increasing, exact worst-case output on U is a*tanh(alpha).
def toy_sweep():
    alpha, eps = 1.0, 0.05
    exact_boundary=(alpha-eps)/math.tanh(alpha)
    rows=[]
    for a in np.linspace(.5,1.35,18):
        x=np.linspace(-alpha,alpha,200001)
        y=a*np.tanh(x)
        violation=float(np.mean(np.abs(y)>alpha-eps))
        worst=float(np.max(np.abs(y)))
        # reachable interval from a dense initial cloud, with empirical diameter
        z=np.linspace(-alpha,alpha,20001)
        diam=[]
        for _ in range(30):
            z=a*np.tanh(z); diam.append(float(z.max()-z.min()))
        
        # Exact symmetric attractor prediction: zero for a<=1; for a>1 solve x=a*tanh(x).
        if a <= 1: predicted_diam=0.0
        else:
            lo,hi=0.0,5.0
            for _ in range(60):
                mid=(lo+hi)/2
                if a*math.tanh(mid)>mid: lo=mid
                else: hi=mid
            predicted_diam=2*((lo+hi)/2)
        rows.append({'a':float(a),'worst':worst,'violation_rate':violation,'diam0':diam[0],'diam30':diam[-1], 'predicted_attractor_diameter':predicted_diam})
    observed=min(r['a'] for r in rows if r['violation_rate']>0)
    # Boundary is grid-resolved; predicted threshold is exact.
    return {'alpha':alpha,'epsilon':eps,'predicted_boundary':exact_boundary,
            'observed_first_violating_grid':observed,'rows':rows,
            'boundary_abs_error':abs(observed-exact_boundary)}

class RNN(nn.Module):
    def __init__(self,h=16):
        super().__init__(); self.h=h
        self.w_in=nn.Linear(1,h); self.w_h=nn.Linear(h,h); self.out=nn.Linear(h,1)
    def transition(self,h,x): return torch.tanh(self.w_in(x)+self.w_h(h))
    def forward(self,x,h=None):
        if h is None: h=torch.zeros(x.shape[0],self.h,device=x.device)
        ys=[]
        for t in range(x.shape[1]):
            h=self.transition(h,x[:,t,:]); ys.append(self.out(h))
        return torch.cat(ys,1).squeeze(-1),h

def mackey(n=1800):
    x=np.zeros(n+31,dtype=np.float32); x[:31]=1.2
    for t in range(30,n+30): x[t+1]=x[t]+.2*x[t-30]/(1+x[t-30]**10)-.1*x[t]
    x=x[30:]; x=(x-x.mean())/(x.std()+1e-6)
    X=np.stack([x[i:i+40] for i in range(n-40)],0)
    Y=np.stack([x[i+1:i+41] for i in range(n-40)],0)
    return torch.tensor(X[:,:,None]),torch.tensor(Y)

def train(trap_lambda, X, Y, steps=500):
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        model=RNN().to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        X,Y=X.to(device),Y.to(device)
        alpha,eps=.8,.05
        for step in range(steps):
            ix=torch.arange((step*32)%len(X),((step*32)%len(X))+32,device=device)%len(X)
            xb,yb=X[ix],Y[ix]
            pred,_=model(xb); loss=F.mse_loss(pred,yb)
            if trap_lambda:
                h=torch.empty(256,model.h,device=device).uniform_(-alpha,alpha)
                z=torch.zeros(256,1,device=device)
                yh=model.transition(h,z)
                trap=F.softplus(torch.abs(yh)-alpha+eps).mean()
                loss=loss+trap_lambda*trap
            opt.zero_grad(); loss.backward(); opt.step()
        return model,device
    except Exception:
        device='cpu'; model=RNN(); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        for step in range(steps):
            ix=torch.arange((step*32)%len(X),((step*32)%len(X))+32)%len(X); xb,yb=X[ix],Y[ix]
            pred,_=model(xb); loss=F.mse_loss(pred,yb)
            if trap_lambda:
                h=torch.empty(256,model.h).uniform_(-.8,.8); yh=model.transition(h,torch.zeros(256,1))
                loss=loss+trap_lambda*F.softplus(torch.abs(yh)-.75).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        return model,device

def eval_model(model,device,X,Y,particles=10000,horizon=1000):
    alpha,eps=.8,.05
    model.eval()
    with torch.no_grad():
        pred,_=model(X[:128].to(device)); mse=float(F.mse_loss(pred,Y[:128].to(device)).cpu())
        h=torch.empty(particles,model.h,device=device).uniform_(-alpha,alpha); zero=torch.zeros(particles,1,device=device)
        violations=[]; maxnorm=[]; clouds=[]
        for t in range(horizon):
            h=model.transition(h,zero); violations.append(float((h.abs()>alpha-eps).any(1).float().mean().cpu()))
            if t in tuple(sorted(set([0, min(9,horizon-1), min(99,horizon-1), horizon-1]))): clouds.append(float((h.max(0).values-h.min(0).values).norm().cpu()))
            maxnorm.append(float(h.norm(dim=1).max().cpu()))
        # one-step empirical certificate on fresh U samples
        cert=float(np.mean(violations[:1]))
        return {'mse':mse,'one_step_violation':cert,'max_norm':max(maxnorm),
                'mean_violation_1000':float(np.mean(violations)),'cloud_diameters':clouds}

def main():
    toy=toy_sweep(); X,Y=mackey()
    
    sweep={}
    for lam in [0., .1, .3, 1., 3., 10.]:
        torch.manual_seed(SEED + int(lam*10))
        m,d=train(lam,X,Y,steps=350)
        sweep[str(lam)]=eval_model(m,d,X,Y,particles=3000,horizon=300)
    results={'toy':toy,'rnn_lambda_sweep':sweep}
    Path('results.json').write_text(json.dumps(results,indent=2))
    print(json.dumps(results,indent=2))
if __name__=='__main__': main()