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

SEED = 471

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def relu_cov(K, variance_scale=2.0):
    # E[ReLU(u) ReLU(v)] for a zero-mean Gaussian pair, followed by He scale.
    d = torch.sqrt(torch.clamp(torch.diag(K), min=1e-12))
    corr = K / (d[:, None] * d[None, :])
    corr = torch.clamp(corr, -1.0, 1.0)
    theta = torch.acos(corr)
    e = (d[:, None] * d[None, :]) * (torch.sin(theta) + (math.pi-theta)*corr) / (2*math.pi)
    return variance_scale * e

def nngp_targets(X, depth):
    K = X @ X.T / X.shape[1]
    Ks = []
    for _ in range(depth):
        K = relu_cov(K)
        Ks.append(K)
    return Ks

class MLP(nn.Module):
    def __init__(self, d, width, depth):
        super().__init__()
        self.layers = nn.ModuleList()
        self.layers.append(nn.Linear(d, width, bias=False))
        for _ in range(depth-1): self.layers.append(nn.Linear(width, width, bias=False))
        self.out = nn.Linear(width, 1, bias=False)
        for m in self.layers:
            nn.init.normal_(m.weight, std=math.sqrt(2/m.in_features))
        nn.init.normal_(self.out.weight, std=1/math.sqrt(width))
    def forward(self, x):
        hs=[]
        h=x
        for layer in self.layers:
            h=torch.relu(layer(h)); hs.append(h)
        return self.out(h), hs

def covariance(h):
    return h @ h.T / h.shape[1]

def train(reg, device, seed):
    seed_all(seed)
    n,d=192,10
    X=torch.randn(n,d,device=device)
    # A nonlinear but learnable fixed task, with a held-out validation set.
    y=(torch.sin(X[:,0]*1.4)+0.35*X[:,1]*X[:,2]).unsqueeze(1)
    Xt=torch.randn(96,d,device=device)
    yt=(torch.sin(Xt[:,0]*1.4)+0.35*Xt[:,1]*Xt[:,2]).unsqueeze(1)
    width,depth=64,3; batch=32
    model=MLP(d,width,depth).to(device)
    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    # The NNGP target is computed from the same input minibatch, and detached.
    history=[]; cov_history=[]
    for step in range(240):
        ix=torch.arange((step*batch)%n, (step*batch)%n+batch, device=device)%n
        xb,yb=X[ix],y[ix]
        pred,hs=model(xb)
        task=((pred-yb)**2).mean()
        targets=nngp_targets(xb,depth)
        covloss=sum(((covariance(h)-k.detach())**2).mean() for h,k in zip(hs,targets))/depth
        loss=task + (0.03*covloss if reg else 0.0)
        opt.zero_grad(); loss.backward(); opt.step()
        if step in (0,39,119,239):
            with torch.no_grad():
                vp,vhs=model(Xt); val=((vp-yt)**2).mean().item()
                # Evaluate deviation on a fresh calibration batch, as proposed.
                _,eh=model(X[:batch]); tk=nngp_targets(X[:batch],depth)
                dev=float(sum(((covariance(h)-k)**2).mean() for h,k in zip(eh,tk))/depth)
            history.append(val); cov_history.append(dev)
    return history[-1], cov_history[-1], history, cov_history

def scaling_check():
    seed_all(SEED)
    d,B=12,24
    X=torch.randn(B,d)
    K=nngp_targets(X,2)[-1]
    widths=[32,64,128,256,512]
    reps=40
    means=[]
    for w in widths:
        errs=[]
        for r in range(reps):
            # independent He network, no output layer; compare postactivation Gram.
            W=torch.randn(d,w)*math.sqrt(2/d)
            h=torch.relu(X@W)
            W2=torch.randn(w,w)*math.sqrt(2/w)
            h=torch.relu(h@W2)
            errs.append(torch.linalg.norm(covariance(h)-K).item())
        means.append(float(np.mean(errs)))
    slope=float(np.polyfit(np.log(widths),np.log(means),1)[0])
    # Also verify the scalar ReLU formula at correlation values by Monte Carlo.
    z=torch.randn(800000,2)
    rho=0.6; z[:,1]=rho*z[:,0]+math.sqrt(1-rho*rho)*z[:,1]
    mc=(torch.relu(z[:,0])*torch.relu(z[:,1])).mean().item()*2
    kk=torch.tensor([[1.,rho],[rho,1.]])
    analytic=relu_cov(kk)[0,1].item()
    return widths,means,slope,mc,analytic

def main():
    try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    except Exception: device=torch.device('cpu')
    try:
        scaling=scaling_check()
        runs=[]
        for s in [471,472,473]:
            b=train(False,device,s); r=train(True,device,s)
            runs.append({'seed':s,'baseline_val':b[0],'idea_val':r[0], 'baseline_cov':b[1], 'idea_cov':r[1], 'baseline_curve':b[2], 'idea_curve':r[2], 'baseline_cov_curve':b[3], 'idea_cov_curve':r[3]})
        out={'device':str(device),'scaling':{'widths':scaling[0],'errors':scaling[1],'loglog_slope':scaling[2],'mc_relu_cov':scaling[3],'analytic_relu_cov':scaling[4]},'runs':runs}
    except Exception as e:
        if device.type=='cuda':
            device=torch.device('cpu'); scaling=scaling_check(); runs=[]
            for s in [471,472,473]:
                b=train(False,device,s); r=train(True,device,s)
                runs.append({'seed':s,'baseline_val':b[0],'idea_val':r[0], 'baseline_cov':b[1], 'idea_cov':r[1]})
            out={'device':'cpu_fallback','scaling':{'widths':scaling[0],'errors':scaling[1],'loglog_slope':scaling[2],'mc_relu_cov':scaling[3],'analytic_relu_cov':scaling[4]},'runs':runs,'cuda_error':repr(e)}
        else: raise
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()