import json, math, os, random
import numpy as np
import torch
from torch import nn
from scipy.integrate import solve_ivp

SEED=2976
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
try:
    if device.type=='cuda': torch.cuda.empty_cache()
except Exception:
    device=torch.device('cpu')

T=0.30; dt=0.03; N=int(T/dt); K=8
# Smooth, positive, strongly variable diffusion on a periodic 1D domain.
def A(x): return 0.20 + 0.15*torch.sin(x)
def An(x): return 0.20 + 0.15*np.sin(x)

def ref_solution(nx=256):
    x=np.linspace(-np.pi,np.pi,nx,endpoint=False); dx=2*np.pi/nx
    # Fourier spectral second derivative, periodic reference solve.
    k=2*np.pi*np.fft.fftfreq(nx,d=dx)
    def rhs(t,v):
        vxx=np.fft.ifft(-(k*k)*np.fft.fft(v)).real
        return -0.5*An(x)*vxx
    sol=solve_ivp(rhs,(T,0),np.cos(2*x),rtol=2e-8,atol=2e-9,method='DOP853')
    return x,sol.y[:,-1]

class MLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,1))
    def forward(self,t,x): return self.net(torch.cat((t,x),1))

def deriv_x(v,x):
    g=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
    h=torch.autograd.grad(g.sum(),x,create_graph=True)[0]
    return g,h

def make_batch(n, mode, model):
    # Time locations are grid points; network learns all time slices jointly.
    ni=torch.randint(0,N,(n,),device=device); t=ni.float().view(-1,1)*dt/T
    x=(torch.rand(n,1,device=device)*2*math.pi-math.pi)/math.pi # normalized coordinate
    # physical coordinate for coefficients; network coordinate is x/pi
    xp=(x*math.pi).detach().requires_grad_(True)
    tn=((ni+1).float().view(-1,1)*dt/T).detach()
    if mode=='raw':
        vn=model(tn,xp/math.pi)
        _,vxx=deriv_x(vn,xp)
        target=vn + dt*0.5*A(xp)*vxx
    else:
        # Frozen kernel is applied with local A0=A(x), using antithetic samples.
        z=torch.randn(n,K//2,device=device); z=torch.cat((z,-z),1)
        xprop=xp.detach()+torch.sqrt(A(xp.detach())*dt)*z
        # periodic wrap; normalized network coordinate remains in [-1,1]
        xwrap=(xprop+math.pi)%(2*math.pi)-math.pi
        tt=tn[:,None,:].expand(n,K,1).reshape(n*K,1)
        xx=xwrap.reshape(n*K,1).detach().requires_grad_(True)
        vv=model(tt,xx/math.pi)
        gx,hx=deriv_x(vv,xx)
        aq=A(xx); a0=A(xp.detach()).repeat_interleave(K).view(-1,1)
        q=0.5*(aq-a0)*hx
        # P V + dt P[(A-A0)Vxx/2]. The local source is zero at patch center,
        # but is retained at propagated points as required by the parametrix formula.
        target=(vv + dt*q).view(n,K,1).mean(1)
    return model(t,x), target.detach()

def train(mode, steps=700):
    torch.manual_seed(SEED+ (0 if mode=='raw' else 1)); model=MLP().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    # Soft terminal anchoring makes the comparison a complete backward solver.
    hist=[]
    for it in range(steps):
        opt.zero_grad(set_to_none=True)
        pred,y=make_batch(96,mode,model)
        loss=((pred-y)**2).mean()
        xt=(torch.rand(96,1,device=device)*2*math.pi-math.pi)/math.pi
        tt=torch.ones_like(xt)
        term=((model(tt,xt)-torch.cos(2*xt*math.pi))**2).mean()
        total=loss+2.0*term
        total.backward(); opt.step()
        if it%50==0: hist.append(float(total.detach().cpu()))
    with torch.no_grad():
        xx=np.linspace(-math.pi,math.pi,256,endpoint=False)
        tx=torch.full((256,1),0.,device=device); xxn=torch.tensor(xx/math.pi,dtype=torch.float32,device=device).view(-1,1)
        pred=model(tx,xxn).cpu().numpy().ravel()
    return model,hist,pred

def main():
    # Core math sanity: Gaussian multiplier for cos(kx), and correction exactly zero for constant A.
    x=torch.linspace(-math.pi,math.pi,10001); k=3.; h=.07; a=.4
    z=torch.randn(300000); mc=torch.cos(k*(x[5000]+math.sqrt(a*h)*z)).mean().item()
    exact=math.cos(k*x[5000].item())*math.exp(-.5*h*a*k*k)
    const_corr=float(torch.max(torch.abs((torch.full_like(x,.4)-.4)*torch.ones_like(x))).item())
    rx,rv=ref_solution()
    results={'device':str(device),'math_check_mc':mc,'math_check_exact':exact,'math_abs_error':abs(mc-exact),'constant_A_correction_max':const_corr,'reference_rms':float(np.sqrt(np.mean(rv*rv)))}
    allhist={}
    for mode in ('raw','parametrix'):
        _,hist,pred=train(mode)
        allhist[mode]=hist
        results[mode+'_final_rel_l2']=float(np.linalg.norm(pred-rv)/np.linalg.norm(rv))
        results[mode+'_loss_trace']=hist
    results['winner_lower_error']='parametrix' if results['parametrix_final_rel_l2']<results['raw_final_rel_l2'] else 'raw'
    with open('results.json','w') as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))

if __name__=='__main__': main()