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

SEED=2852
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception: device=torch.device('cpu')

# periodic y interpolation; x diffusion is a normalized Gaussian depthwise convolution
def diffusion_x(f, h, dx):
    # f [B,C,Nx,Ny], Gaussian variance 2h
    sigma=math.sqrt(2*h)/dx
    r=max(1,int(math.ceil(3*sigma)))
    z=torch.arange(-r,r+1,device=f.device,dtype=f.dtype)
    w=torch.exp(-0.5*(z/sigma)**2); w=w/w.sum()
    out=torch.zeros_like(f)
    for i,wi in enumerate(w): out += wi*torch.roll(f, int(i-r), dims=2)
    return out

def transport_y(f,h,x,ymin=-math.pi,ymax=math.pi):
    # source f(x, y+h*x); periodic linear interpolation
    ny=f.shape[-1]; dy=(ymax-ymin)/ny
    shift=h*x[:,None] # [Nx,1]
    yy=(torch.arange(ny,device=f.device,dtype=f.dtype)[None,:]*dy+ymin+shift)
    q=(yy-ymin)/dy % ny
    j0=torch.floor(q).long(); a=(q-j0).to(f.dtype); j1=(j0+1)%ny
    return f.gather(3,j0[None,None].expand(f.shape[0],f.shape[1],-1,-1))*(1-a)[None,None] + f.gather(3,j1[None,None].expand(f.shape[0],f.shape[1],-1,-1))*a[None,None]

def kinetic(f,h,x,dx): return transport_y(diffusion_x(f,h,dx),h,x)

class KineticLayer(nn.Module):
    def __init__(self,c,h,x,dx):
        super().__init__(); self.h=h; self.register_buffer('x',x); self.dx=dx
        self.gate=nn.Conv2d(c,c,1); nn.init.constant_(self.gate.weight,0); nn.init.constant_(self.gate.bias,1.5)
        self.mix=nn.Conv2d(c,c,1)
    def forward(self,f):
        z=kinetic(f,self.h,self.x,self.dx)
        g=torch.sigmoid(self.gate(f)); return self.mix(f+g*(z-f))
class LocalBaseline(nn.Module):
    def __init__(self,c):
        super().__init__(); self.net=nn.Sequential(nn.Conv2d(c,c,3,padding=1,padding_mode='circular'),nn.GELU(),nn.Conv2d(c,c,1))
    def forward(self,f): return self.net(f)

def make_fields(n,c,nx,ny,device):
    # Smooth random phase-space fields, with enough variation to expose y transport.
    q=torch.randn(n,c,nx,ny,device=device)
    for _ in range(3): q=(q+torch.roll(q,1,2)+torch.roll(q,-1,2)+torch.roll(q,1,3)+torch.roll(q,-1,3))/5
    return q

def math_checks():
    nx,ny=64,128; xmin,xmax=-2,2; dx=(xmax-xmin)/nx
    x=torch.arange(nx,dtype=torch.float64)*dx+xmin; yy=torch.arange(ny,dtype=torch.float64)*(2*math.pi/ny)-math.pi
    # diffusion Fourier mode prediction: exp(-h*k^2)
    k=3.0; xx=x[:,None]; f=torch.cos(k*xx).expand(1,1,nx,ny).clone()
    hs=np.array([.01,.025,.05,.1,.2]); ratios=[]
    for h in hs:
        d=diffusion_x(f,h,dx); ratios.append((d.abs().mean()/f.abs().mean()).item())
    pred=np.exp(-hs*k*k); rel=float(np.max(np.abs(np.array(ratios)-pred)/pred))
    # transport preserves L2; convex gated update is nonexpansive for constant g against zero
    rng=torch.Generator().manual_seed(SEED); r=torch.randn(1,1,nx,ny,generator=rng,dtype=torch.float64)
    norms=[]; disp=[]
    for h in [.02,.05,.1,.2]:
        z=kinetic(r,h,x,dx); norms.append((z.norm()/r.norm()).item())
        # characteristic displacement RMS (on physical coordinates)
        disp.append(h*float(torch.sqrt(torch.mean(x*x))))
    # contraction test for g=0,.25,.5,1 on a random vector; diffusion+shift operator is empirically contractive
    contractions=[]
    for g in [0,.25,.5,1]:
        z=r+g*(kinetic(r,.1,x,dx)-r); contractions.append((z.norm()/r.norm()).item())
    return {'diffusion_h':hs.tolist(),'diffusion_observed':ratios,'diffusion_predicted':pred.tolist(),'diffusion_max_relative_error':rel,
            'kinetic_norm_ratios_h_[.02,.05,.1,.2]':norms,'transport_rms_displacement':disp,
            'displacement_over_h':(np.array(disp)/np.array([.02,.05,.1,.2])).tolist(),
            'gated_norm_ratios_g_[0,.25,.5,1]':contractions}

def train_compare():
    c,nx,ny=2,24,48; h=.12; xmin,xmax=-2,2; dx=(xmax-xmin)/nx
    x=torch.arange(nx,device=device)*dx+xmin
    train=make_fields(96,c,nx,ny,device); val=make_fields(32,c,nx,ny,device)
    with torch.no_grad(): target=kinetic(train,h,x,dx); vtarget=kinetic(val,h,x,dx)
    models={'baseline':LocalBaseline(c).to(device),'idea':KineticLayer(c,h,x,dx).to(device)}
    results={}
    for name,m in models.items():
        opt=torch.optim.Adam(m.parameters(),lr=3e-3)
        for step in range(350):
            idx=torch.randint(0,len(train),(16,),device=device); out=m(train[idx]); loss=F.mse_loss(out,target[idx])
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=m(val); mse=F.mse_loss(pred,vtarget).item()
            # rollout by repeatedly applying learned transition, measured against true operator
            a=val.clone(); b=val.clone(); errs_a=[]; errs_b=[]
            for _ in range(5):
                a=m(a); b=kinetic(b,h,x,dx); errs_a.append(F.mse_loss(a,b).item())
            results[name]={'val_one_step_mse':mse,'rollout_mse_steps_1_to_5':errs_a,'parameters':sum(p.numel() for p in m.parameters())}
    return results

if __name__=='__main__':
    checks=math_checks()
    try:
        comparison=train_compare(); used_device=str(device)
    except Exception as e:
        print('CUDA/backend failure, falling back to CPU:', repr(e))
        device=torch.device('cpu'); torch.manual_seed(SEED)
        comparison=train_compare(); used_device='cpu (fallback)'
    print(json.dumps({'device':used_device,'math_checks':checks,'comparison':comparison},indent=2))