import json, math, random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.special import iv, jv
from scipy.optimize import brentq
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

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

class BernsteinResolvent(nn.Module):
    def __init__(self, K=8, tau=0.5, eps=1e-4):
        super().__init__(); self.tau=tau; self.eps=eps
        poles=torch.exp(torch.linspace(math.log(.03), math.log(30.), K))
        self.ua=nn.Parameter(torch.tensor(math.log(math.expm1(.2))))
        self.uc=nn.Parameter(torch.full((K,), math.log(math.expm1(.15))))
        self.register_buffer('poles', poles)
    def forward(self,x):
        y=self.eps+F.relu(x); s=y.pow(self.tau)
        a=F.softplus(self.ua); c=F.softplus(self.uc)
        return a*s+(c*s[...,None]/(s[...,None]+self.poles)).sum(dim=-1)

class MLP(nn.Module):
    def __init__(self,kind='gelu',K=8,tau=.5):
        super().__init__(); self.kind=kind
        self.layers=nn.ModuleList([nn.Linear(64,128),nn.Linear(128,128),nn.Linear(128,128)])
        self.out=nn.Linear(128,10)
        self.act=BernsteinResolvent(K,tau) if kind=='resolvent' else nn.GELU()
    def forward(self,x):
        for layer in self.layers: x=self.act(layer(x))
        return self.out(x)

def math_checks():
    # Check W_nu partial fraction against the Bessel ratio at positive s.
    nu=.3; ss=np.array([.01,.1,1.,10.])
    # Locate positive zeros for the fractional-order Bessel function J_{nu+1}.
    order=nu+1; zeros=[]
    step=math.pi/8
    left=1e-7; fl=jv(order,left)
    while len(zeros)<500:
        right=left+step; fr=jv(order,right)
        if fl*fr<0: zeros.append(brentq(lambda z:jv(order,z),left,right))
        left,fl=right,fr
    lam=np.asarray(zeros)**2
    approx=2*(nu+1)+2*np.array([np.sum(s/(s+lam)) for s in ss])
    exact=np.array([math.sqrt(s)*iv(nu,math.sqrt(s))/iv(nu+1,math.sqrt(s)) for s in ss])
    relerr=float(np.max(np.abs(approx-exact)/np.abs(exact)))
    # Dense autograd derivative signs for finite analogue on x>=0.
    x=torch.linspace(0,100,4000,requires_grad=True)
    vals=BernsteinResolvent(K=8,tau=.5)(x[:,None]).squeeze()
    d=torch.autograd.grad(vals.sum(),x,create_graph=True)[0]
    d2=torch.autograd.grad(d.sum(),x)[0]
    mono=float(d.min()); conc=float(d2.max())
    # Deliberately outside range: numerical concavity check for tau=.75.
    x2=torch.linspace(1e-3,100,4000,requires_grad=True)
    v2=BernsteinResolvent(K=8,tau=.75)(x2[:,None]).squeeze()
    d21=torch.autograd.grad(v2.sum(),x2,create_graph=True)[0]
    d22=torch.autograd.grad(d21.sum(),x2)[0]
    return {'bessel_partial_fraction_max_relative_error':relerr,
            'tau_0.5_min_first_derivative':mono,'tau_0.5_max_second_derivative':conc,
            'tau_0.75_max_second_derivative':float(d22.max())}

def train(kind,tau=.5,steps=500):
    data=load_digits(); X=data.data.astype('float32'); y=data.target
    X=StandardScaler().fit_transform(X).astype('float32')
    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
    xt=torch.tensor(Xtr,device=device); yt=torch.tensor(ytr,device=device)
    xv=torch.tensor(Xte,device=device); yv=torch.tensor(yte,device=device)
    torch.manual_seed(SEED)
    model=MLP(kind,tau=tau).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    losses=[]; gradnorms=[]; bs=96; g=torch.Generator(device=device).manual_seed(SEED)
    for step in range(steps):
        ind=torch.randint(0,len(xt),(bs,),generator=g,device=device)
        opt.zero_grad(set_to_none=True); loss=F.cross_entropy(model(xt[ind]),yt[ind]); loss.backward()
        gn=float(torch.sqrt(sum((p.grad.detach()**2).sum() for p in model.parameters() if p.grad is not None)).cpu())
        gradnorms.append(gn); losses.append(float(loss.detach().cpu())); opt.step()
    with torch.no_grad():
        pred=model(xv).argmax(1); acc=float((pred==yv).float().mean().cpu())
    return {'final_train_loss':losses[-1],'test_accuracy':acc,
            'mean_grad_norm':float(np.mean(gradnorms)),'grad_norm_std':float(np.std(gradnorms)),
            'max_grad_norm':float(np.max(gradnorms))}

def main():
    result={'seed':SEED,'device':device,'math':math_checks(),'experiment':{}}
    for kind,tau in [('gelu',.5),('resolvent',.5),('resolvent',.75)]:
        key=kind if kind=='gelu' else kind+'_tau_'+str(tau).replace('.','_')
        result['experiment'][key]=train(kind,tau)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()