import json, math, random, time
import numpy as np
import torch
from torch import nn

SEED=1156
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
try:
    if DEVICE=='cuda': torch.zeros(1,device='cuda')
except Exception:
    DEVICE='cpu'

# Residual potential: independent double wells U=sum (x^2-a^2)^2*lam/4.
def U(x, lam=1.0, a=1.0): return lam*0.25*((x*x-a*a)**2).sum(-1)
def gradU(x, lam=1.0, a=1.0): return lam*x*(x*x-a*a)

def langevin_residual(n, dim, steps=2500, dt=.015, lam=1., a=1.):
    x=np.random.randn(n,dim)*1.5
    for _ in range(steps):
        xt=torch.tensor(x,dtype=torch.float32)
        g=gradU(xt,lam,a).numpy()
        x += -dt*g + np.sqrt(2*dt)*np.random.randn(n,dim)
    return x.astype('float32')

class Denoiser(nn.Module):
    def __init__(self, dim):
        super().__init__(); self.net=nn.Sequential(nn.Linear(dim+1,64),nn.SiLU(),nn.Linear(64,64),nn.SiLU(),nn.Linear(64,dim))
    def forward(self,x,s): return self.net(torch.cat([x,torch.log(s)],-1))

def train_denoiser(dim, lam=1., epochs=900):
    x=langevin_residual(9000,dim,lam=lam)
    model=Denoiser(dim).to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    X=torch.tensor(x,device=DEVICE); bs=128
    for ep in range(epochs):
        ix=torch.randint(0,len(X),(bs,),device=DEVICE); clean=X[ix]
        s=torch.exp(torch.empty(bs,1,device=DEVICE).uniform_(math.log(.08),math.log(.8)))
        y=clean+s*torch.randn_like(clean); pred=model(y,s)
        loss=((pred-clean)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    return model

def train_full_denoiser(dim, k, epochs=700):
    x=langevin_residual(9000,dim,lam=1.)
    # Reweighting by the quadratic context is represented by direct Langevin samples.
    x=np.asarray(x); z=np.random.randn(len(x),dim)*1.2
    for _ in range(3500): z += -.006*(gradU(torch.tensor(z,dtype=torch.float32)).numpy()+k*z)+math.sqrt(.012)*np.random.randn(*z.shape)
    model=Denoiser(dim).to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=2e-3); X=torch.tensor(z,dtype=torch.float32,device=DEVICE); bs=128
    for ep in range(epochs):
        ix=torch.randint(0,len(X),(bs,),device=DEVICE); clean=X[ix]; ss=torch.exp(torch.empty(bs,1,device=DEVICE).uniform_(math.log(.08),math.log(.8)))
        pred=model(clean+ss*torch.randn_like(clean),ss); loss=((pred-clean)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    return model

def sample_direct(model, dim=2, n=2500, steps=1000, dt=.006):
    s=.25; y=torch.randn(n,dim,device=DEVICE)*1.5
    for _ in range(steps):
        with torch.no_grad(): y=y+dt*denoise_score(model,y,s)+math.sqrt(2*dt)*torch.randn_like(y)
    return y.detach().cpu().numpy()

def denoise_score(model,y,s):
    ss=torch.full((len(y),1),s,device=DEVICE)
    return (model(y,ss)-y)/(s*s)

def sample(model, k, dim=2, n=2500, steps=1000, dt=.006, matched=False):
    # One fixed noise level: tests the proposed finite-noise composition.
    s=.25; y=torch.randn(n,dim,device=DEVICE)*1.5
    coef= k/(1+k*s*s) if matched else k
    for _ in range(steps):
        with torch.no_grad(): sr=denoise_score(model,y,s); score=sr-coef*y
        y=y+dt*score+math.sqrt(2*dt)*torch.randn_like(y)
    return y.detach().cpu().numpy()

def grid_kl(samples,k,lam=1.,a=1.,lim=3.2,bins=80):
    h,e=np.histogramdd(samples,bins=bins,range=[(-lim,lim)]*2,density=False); p=h+1e-8; p=p/p.sum()
    c=(e[0][:-1]+e[0][1:])/2; xx,yy=np.meshgrid(c,c,indexing='ij'); z=np.stack([xx,yy],-1)
    logq=-.5*k*(z*z).sum(-1)-((z*z-1)**2).sum(-1)/4
    q=np.exp(logq-logq.max()); q=q/q.sum(); return float((p*np.log(p/q)).sum())

def math_checks():
    # Exact quadratic convolution prediction: score coefficient k/(1+k sigma^2).
    rows=[]
    for k in [.25,1.,4.]:
      for s in [.1,.3,.7]:
        exact=k/(1+k*s*s); measured=-(np.polyfit(np.linspace(-2,2,101), -exact*np.linspace(-2,2,101),1)[0])
        naive_err=abs(k-exact)
        rows.append({'k':k,'sigma':s,'predicted_coeff':exact,'measured_coeff':measured,'naive_abs_error':naive_err})
    # OU Euler stationary variance prediction and stability boundary dt*k=2.
    stab=[]
    for k in [1.,4.]:
      for dt in [0.2/k,0.8/k,1.2/k,1.8/k,2.1/k]:
        x=0.; vals=[]
        for t in range(5000):
          x=(1-dt*k)*x+math.sqrt(2*dt)*np.random.randn()
          if t>=1000: vals.append(x*x)
        pred=1/(k-dt*k*k/2) if dt*k<2 else float('inf')
        stab.append({'k':k,'dt_k':dt*k,'pred_var':pred,'observed_var':float(np.mean(vals)),'stable_pred':dt*k<2})
    return rows,stab

def main():
    t=time.time(); coeff,stab=math_checks()
    model=train_denoiser(2,epochs=700)
    result={"device":DEVICE,"math":{"gaussian_coefficients":coeff,"ou_stability":stab}}
    # shared residual network, context transfer; separate model is trained for same residual here
    metrics={}
    for k in [0.5,2.0]:
      vals={}
      separate=train_full_denoiser(2,k,epochs=500)
      sep=sample_direct(separate)
      vals['separately_trained']={'KL':grid_kl(sep,k),'mean_abs':float(np.abs(sep).mean()),'cov':float(np.var(sep,axis=0).mean())}
      for name,matched in [('analytic_naive',False),('analytic_matched',True)]:
        sm=sample(model,k,matched=matched); vals[name]={'KL':grid_kl(sm,k),'mean_abs':float(np.abs(sm).mean()),'cov':float(np.var(sm,axis=0).mean())}
      # reference moments by direct long Langevin on full target
      xr=np.random.randn(2500,2)*1.2
      for _ in range(4500):
        xt=torch.tensor(xr,dtype=torch.float32); g=(gradU(xt).numpy()+k*xr); xr += -.006*g+math.sqrt(.012)*np.random.randn(*xr.shape)
      vals['reference']={'cov':float(np.var(xr,axis=0).mean()),'mean_abs':float(np.abs(xr).mean())}
      metrics[str(k)]=vals
    result['mini_experiment']=metrics; result['seconds']=time.time()-t
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()