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

SEED = 1358
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
M = 1.0

# Klein-Gordon reduced to a single Fourier mode: u_t=v, v_t=-m sin(u).
def F(z):
    u, v = z[..., 0], z[..., 1]
    return np.stack([v, -M*np.sin(u)], axis=-1)

def rk4(z, dt):
    k1=F(z); k2=F(z+.5*dt*k1); k3=F(z+.5*dt*k2); k4=F(z+dt*k3)
    return z + dt*(k1+2*k2+2*k3+k4)/6

def math_checks():
    bg=np.array([.73,-.21]); r=np.array([.18,-.27]); bgdot=np.array([bg[1],-.11])
    defect=F(bg)-bgdot; ru=np.sin(bg[0]+r[0])-np.sin(bg[0])-np.cos(bg[0])*r[0]
    rhs=np.array([r[1],-M*r[0]+defect[1]-M*(np.cos(bg[0])-1)*r[0]-M*ru])
    identity_err=float(np.max(np.abs(F(bg+r)-bgdot-rhs)))
    eps=np.logspace(-5,-1,9)
    rem=np.array([abs(np.sin(bg[0]+e)-np.sin(bg[0])-np.cos(bg[0])*e) for e in eps])
    slope=float(np.polyfit(np.log(eps),np.log(rem),1)[0])
    coeff=float(np.median(rem[:5]/eps[:5]**2)); coeff_pred=abs(np.sin(bg[0]))/2
    # sigmoid gate g=1/(1+exp(-a(log d-b))); midpoint predicted d=exp(b)
    a=3.; b=math.log(.08); ds=np.logspace(-3,0,401)
    gates=1/(1+np.exp(-a*(np.log(ds)-b))); d_mid=float(ds[np.argmin(abs(gates-.5))])
    # logit(g) versus log(d) slope predicted a
    fit=float(np.polyfit(np.log(ds),np.log(gates/(1-gates)),1)[0])
    return {'identity_max_error':identity_err,'remainder_loglog_slope':slope,
            'remainder_quadratic_slope_predicted':2.,'remainder_coeff_est':coeff,
            'remainder_coeff_predicted':coeff_pred,'gate_midpoint_observed':d_mid,
            'gate_midpoint_predicted':math.exp(b),'gate_logit_slope_observed':fit,
            'gate_logit_slope_predicted':a}

def make_data(n=1800, steps=25, dt=.04):
    x=np.random.uniform(-1.4,1.4,(n,2)); ys=[]
    for _ in range(steps):
        y=rk4(x,dt); ys.append((x.copy(),y.copy())); x=y
    return np.concatenate([a for a,b in ys]),np.concatenate([b for a,b in ys])

class Net(nn.Module):
    def __init__(self, inp, out=2):
        super().__init__(); self.net=nn.Sequential(nn.Linear(inp,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,out))
    def forward(self,x): return self.net(x)

def train(kind, X, Y, steps=500):
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        tx=torch.tensor(X,dtype=torch.float32,device=dev); ty=torch.tensor(Y,dtype=torch.float32,device=dev)
        net=Net(2 if kind=='full' else 4).to(dev); opt=torch.optim.Adam(net.parameters(),lr=2e-3)
        for i in range(steps):
            idx=torch.randint(0,len(tx),(128,),device=dev); z=tx[idx]; target=ty[idx]
            if kind=='full': pred=net(z)
            else:
                # cheap background is a scaled oscillator drift; defect is explicitly supplied
                bg=.82*z; r=z-bg; bgdot=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1)*.82
                fbg=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1); defect=fbg-bgdot
                inp=torch.cat((bg,r,defect[:,1:2]),1) # 5 features
                if net.net[0].in_features != 5: net=Net(5).to(dev); opt=torch.optim.Adam(net.parameters(),lr=2e-3)
                pred=bg + dt_step(z,bg,r,defect,net(inp))
            loss=((pred-target)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            z=tx[:512]; target=ty[:512]
            if kind=='full': pred=net(z)
            else:
                bg=.82*z; r=z-bg; bgdot=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1)*.82
                fbg=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1); defect=fbg-bgdot
                inp=torch.cat((bg,r,defect[:,1:2]),1); pred=bg+dt_step(z,bg,r,defect,net(inp))
            one=float(torch.mean((pred-target)**2).sqrt().cpu())
        return net,one,dev
    except Exception as e:
        # CPU retry is intentionally simple and deterministic.
        if dev!='cpu': torch.cuda.empty_cache(); torch.cuda.is_available=lambda:False; return train(kind,X,Y,steps)
        raise e

def dt_step(z,bg,r,defect,closure,dt=.04):
    linear=torch.stack((r[:,1],-r[:,0]),1)
    corr=torch.stack((torch.zeros_like(r[:,0]),-(torch.cos(bg[:,0])-1)*r[:,0]),1)
    gate=torch.sigmoid(3.0*(torch.log(torch.linalg.vector_norm(defect,dim=1)+1e-6)-math.log(.08))).unsqueeze(1)
    return dt*(linear+defect+corr+gate*closure)

def rollout(kind, net, init, horizon=80, dt=.04):
    z=torch.tensor(init,dtype=torch.float32,device=next(net.parameters()).device); out=[]
    with torch.no_grad():
        for _ in range(horizon):
            if kind=='full': z=net(z)
            else:
                bg=.82*z; r=z-bg; bgdot=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1)*.82
                fbg=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1); defect=fbg-bgdot
                inp=torch.cat((bg,r,defect[:,1:2]),1); z=bg+dt_step(z,bg,r,defect,net(inp))
            out.append(z.cpu().numpy())
    return np.stack(out)

def main():
    checks=math_checks(); X,Y=make_data(); split=int(.8*len(X));
    results={}
    for k in ('full','residual'):
        net,one,dev=train(k,X[:split],Y[:split]); pred=rollout(k,net,X[split:split+64],60)
        true=[]; z=X[split:split+64].copy()
        for _ in range(60): z=rk4(z,.04); true.append(z.copy())
        longerr=float(np.sqrt(np.mean((pred-np.array(true))**2)))
        results[k]={'one_step_rmse':one,'rollout_rmse':longerr,'device':dev}
    Path('results.json').write_text(json.dumps({'math_checks':checks,'comparison':results},indent=2))
    print(json.dumps({'math_checks':checks,'comparison':results},indent=2))

if __name__=='__main__': main()