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

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

# ---- exact finite empirical mean-field layer ----
class ParticleNet(nn.Module):
    def __init__(self,d,m):
        super().__init__(); self.m=m; self.d=d
        self.a=nn.Parameter(0.25*torch.randn(m, device=device))
        self.w=nn.Parameter(torch.randn(m,d, device=device)/math.sqrt(d))
    def forward(self,x):
        return (F.relu(x @ self.w.T)*self.a).mean(1,keepdim=True)
    def project(self,A=3.0):
        with torch.no_grad():
            self.a.clamp_(-A,A)

class DenseNet(nn.Module):
    def __init__(self,d,m):
        super().__init__(); self.w=nn.Parameter(torch.randn(m,d,device=device)/math.sqrt(d)); self.a=nn.Parameter(0.25*torch.randn(m,device=device))
    def forward(self,x): return (F.relu(x@self.w.T)*self.a).mean(1,keepdim=True)

def math_check():
    torch.manual_seed(SEED+1)
    d,m,n=5,17,23
    a=torch.randn(m); w=torch.randn(m,d); x=torch.randn(n,d)
    lhs=(F.relu(x@w.T)*a).mean(1)
    rhs=torch.stack([(a*F.relu(w@x[i])).mean() for i in range(n)])
    pred_err=(lhs-rhs).abs().max().item()
    # Langevin variance: increments with zero gradient have variance 2 eta lambda.
    eta,lam=0.013,0.7; reps=20000
    z=torch.randn(reps)*math.sqrt(2*eta*lam)
    var=z.var(unbiased=True).item(); target=2*eta*lam
    return {'predictor_max_abs_error':pred_err,'noise_variance':var,'noise_target':target,'relative_variance_error':abs(var-target)/target}

def make_data(d=20,ntrain=2400,ntest=1200):
    g=torch.Generator(device=device).manual_seed(SEED+4)
    # two distinct teacher directions: multi-index test of retained diversity
    q,_=torch.linalg.qr(torch.randn(d,2,device=device,generator=g))
    Xtr=torch.randn(ntrain,d,device=device,generator=g); Xte=torch.randn(ntest,d,device=device,generator=g)
    def target(x):
        return (0.9*F.relu(x@q[:,0])+0.7*F.relu(x@q[:,1])-0.25*F.relu(-x@q[:,0])).unsqueeze(1)
    ytr=target(Xtr); yte=target(Xte)
    return Xtr,ytr,Xte,yte,q

def alignment(model,q):
    # max absolute cosine to either teacher direction, and fraction of particles assigned each mode
    wn=F.normalize(model.w,dim=1); qn=F.normalize(q,dim=0)
    c=(wn@qn).abs(); best=c.max(1).values
    assign=c.argmax(1)
    return float(best.mean()), float((best>0.7).float().mean()), [int((assign==k).sum()) for k in range(2)]

def train(kind,X,y,Xte,yte,q,m=64,steps=700,batch=96):
    torch.manual_seed(SEED+9)
    model=ParticleNet(X.shape[1],m) if kind=='particle' else DenseNet(X.shape[1],m)
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=0.025)
    losses=[]; t0=time.time(); lam_hi,lam_lo=1.0,0.01; anneal=steps
    for t in range(steps):
        ix=torch.randint(0,X.shape[0],(batch,),device=device)
        loss=F.mse_loss(model(X[ix]),y[ix]); opt.zero_grad(); loss.backward()
        # Adam is used for fair neural optimization; particle dynamics adds the specified diffusion.
        opt.step()
        if kind=='particle':
            lam=lam_lo+0.5*(lam_hi-lam_lo)*(1+math.cos(math.pi*t/anneal))
            eta=0.025
            with torch.no_grad():
                model.a.add_(math.sqrt(2*eta*lam)*torch.randn_like(model.a))
                model.w.add_(math.sqrt(2*eta*lam)*torch.randn_like(model.w))
                model.project()
        if t in [0,99,299,699]:
            with torch.no_grad(): losses.append(float(F.mse_loss(model(Xte),yte)))
    with torch.no_grad(): test=float(F.mse_loss(model(Xte),yte)); trainloss=float(F.mse_loss(model(X),y))
    al=alignment(model,q)
    diversity=float(torch.pdist(F.normalize(model.w,dim=1)).mean())
    return {'test_mse':test,'train_mse':trainloss,'checkpoints':losses,'alignment_mean':al[0],'aligned_fraction':al[1],'mode_counts':al[2],'weight_diversity':diversity,'seconds':time.time()-t0}

def main():
    check=math_check(); X,y,Xte,yte,q=make_data()
    base=train('baseline',X,y,Xte,yte,q); idea=train('particle',X,y,Xte,yte,q)
    out={'device':str(device),'seed':SEED,'math_check':check,'baseline':base,'particle':idea,'config':{'width':64,'dimension':20,'steps':700,'batch':96,'schedule':'lambda 1.0 -> 0.01 cosine','teacher':'two-index ReLU regression'}}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()