import json, math, random, time
from pathlib import Path
import numpy as np

SEED = 1234
np.random.seed(SEED)
random.seed(SEED)


def h_np(z, eta, nu, sigma, kmin=-40, kmax=40):
    z = np.asarray(z, dtype=np.float64)
    kk = np.ceil(-np.log(z) / np.log(sigma)).astype(np.int64)
    kk = np.clip(kk, kmin, kmax)
    sk = sigma ** kk
    y = ((eta-nu)/(sigma-1.0))*sk*z + (-eta+sigma*nu)/(sigma-1.0)
    return y, kk


def mechanism_checks():
    rows = []
    for sigma in [1.5, 2.0, 3.0, 4.0]:
        eta, nu = 1.7, 0.35
        # Prediction 1: k changes at b_k=sigma^{-k}.
        boundary_err = []
        boundary_fail = 0
        for k in range(-5, 6):
            b = sigma ** (-k)
            kl = h_np([b*(1-1e-10)], eta, nu, sigma)[1][0]
            kr = h_np([b*(1+1e-10)], eta, nu, sigma)[1][0]
            if not (kl == k+1 and kr == k):
                boundary_fail += 1
            boundary_err.append(abs((sigma**(-k))-b)/b)
        # Prediction 2: slope ratio between neighboring branches is sigma.
        slopes=[]
        for k in range(-4, 5):
            z = sigma**(-k) * math.sqrt(1.0/sigma)
            d = z*1e-6
            yp = h_np([z+d], eta, nu, sigma)[0][0]
            ym = h_np([z-d], eta, nu, sigma)[0][0]
            slopes.append((yp-ym)/(2*d))
        ratios=np.array(slopes[1:])/np.array(slopes[:-1])
        slope_err=float(np.max(np.abs(ratios-sigma)))
        # Prediction 3: equivariance error is roundoff, when clipping is inactive.
        z=np.logspace(-3,3,2000)
        y,_=h_np(z,eta,nu,sigma,kmin=-100,kmax=100)
        ys,_=h_np(z/sigma,eta/sigma,nu/sigma,sigma,kmin=-100,kmax=100)
        eq_rel=float(np.max(np.abs(ys-y/sigma)/(1e-12+np.abs(y/sigma))))
        rows.append({"sigma":sigma,
                     "boundary_relative_error":float(max(boundary_err)),
                     "boundary_failures":boundary_fail,
                     "slope_ratio_observed_mean":float(np.mean(ratios)),
                     "slope_ratio_observed_min":float(np.min(ratios)),
                     "slope_ratio_observed_max":float(np.max(ratios)),
                     "slope_ratio_predicted":sigma,
                     "slope_ratio_max_abs_error":slope_err,
                     "equivariance_max_relative_error":eq_rel})
    return rows


def signed_log_activation_torch(x, eta=1.0, nu=0.0, sigma=2.0, eps=1e-6,
                                kmin=-12, kmax=12):
    import torch
    z=torch.abs(x)+eps
    # Straight-through k: forward integer bins, derivative zero through k.
    k=torch.ceil(-torch.log(z)/math.log(sigma)).clamp(kmin,kmax)
    sk=torch.exp(k*math.log(sigma))
    y=((eta-nu)/(sigma-1))*sk*z + (-eta+sigma*nu)/(sigma-1)
    return torch.sign(x)*y


def mlp_experiment():
    import torch
    from sklearn.datasets import load_digits
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler
    torch.manual_seed(SEED); np.random.seed(SEED)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        if device == "cuda": torch.cuda.empty_cache()
    except Exception:
        device="cpu"
    d=load_digits()
    X=d.data.astype(np.float32); y=d.target.astype(np.int64)
    X=StandardScaler().fit_transform(X).astype(np.float32)
    Xtr,Xv,ytr,yv=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)
    Xval=torch.tensor(Xv,device=device); yval=torch.tensor(yv,device=device)
    class Net(torch.nn.Module):
        def __init__(self, kind):
            super().__init__(); self.kind=kind
            self.l1=torch.nn.Linear(64,96); self.l2=torch.nn.Linear(96,96); self.l3=torch.nn.Linear(96,10)
        def forward(self,x):
            x=self.l1(x)
            x=torch.nn.functional.relu(x) if self.kind=='relu' else signed_log_activation_torch(x,eta=1.,nu=0.,sigma=2.)
            x=self.l2(x)
            x=torch.nn.functional.relu(x) if self.kind=='relu' else signed_log_activation_torch(x,eta=1.,nu=0.,sigma=2.)
            return self.l3(x)
    out={}
    for kind in ['relu','logscale']:
        torch.manual_seed(SEED); net=Net(kind).to(device)
        opt=torch.optim.Adam(net.parameters(),lr=2e-3)
        t0=time.time()
        for step in range(400):
            idx=torch.randint(0,len(Xt),(64,),device=device)
            loss=torch.nn.functional.cross_entropy(net(Xt[idx]),yt[idx])
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            logits=net(Xval); acc=(logits.argmax(1)==yval).float().mean().item()
            scaled=[]
            for q in [-2,-1,0,1,2]:
                scaled.append((net(Xval*(2.0**q)).argmax(1)==yval).float().mean().item())
        out[kind]={"validation_accuracy":acc,"rescaling_accuracy_q_-2..2":scaled,
                   "mean_scaled_accuracy":float(np.mean(scaled)),"seconds":time.time()-t0}
    return {"device":device,"results":out}

if __name__ == '__main__':
    result={"seed":SEED,"mechanism_checks":mechanism_checks()}
    try:
        result["mlp"] = mlp_experiment()
    except Exception as e:
        result["mlp_error"] = repr(e)
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))