import json, math, random
import numpy as np
import torch
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from scipy.stats import spearmanr

SEED = 3016
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(2, 48), torch.nn.ReLU(), torch.nn.Dropout(.35),
            torch.nn.Linear(48, 48), torch.nn.ReLU(), torch.nn.Dropout(.35),
            torch.nn.Linear(48, 2))
    def forward(self, x): return self.net(x)

def math_sanity():
    # Empirical Phi uses local ball mass. A sparse cluster should have larger Phi.
    x = np.r_[np.linspace(-1, 1, 31), np.linspace(5, 5.12, 5)]
    d = np.abs(x[:, None] - x[None, :]); anchor = 15
    def phi(i):
        v = d[i, anchor]
        rs = np.linspace(max(1e-5, 4*v/500), 4*v, 500)
        masses = np.array([(d[i] <= r).mean() for r in rs])
        return np.trapz(np.log(1/masses), rs)
    ph = np.array([phi(i) for i in range(len(x))])
    return {"dense_cluster_mean_phi": float(ph[:31].mean()),
            "sparse_cluster_mean_phi": float(ph[31:].mean()),
            "sparse_larger": bool(ph[31:].mean() > ph[:31].mean()),
            "integrand_monotonicity": bool(math.log(1/.05) > math.log(1/.5))}

def stochastic_logits(model, X, K):
    model.train()  # activate dropout, without gradients
    vals=[]
    with torch.no_grad():
        xt=torch.tensor(X, dtype=torch.float32)
        for _ in range(K): vals.append(model(xt).cpu().numpy())
    return np.stack(vals)

def main():
    sanity=math_sanity()
    X,y=make_moons(n_samples=1800, noise=.25, random_state=SEED)
    X=(X-X.mean(0))/(X.std(0)+1e-8)
    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.35,random_state=SEED,stratify=y)
    model=Net(); opt=torch.optim.Adam(model.parameters(),lr=2e-3,weight_decay=1e-4)
    tx=torch.tensor(Xtr,dtype=torch.float32); ty=torch.tensor(ytr,dtype=torch.long)
    model.train()
    for epoch in range(180):
        perm=torch.randperm(len(tx))
        for j in range(0,len(tx),128):
            ix=perm[j:j+128]; loss=torch.nn.functional.cross_entropy(model(tx[ix]),ty[ix])
            opt.zero_grad(); loss.backward(); opt.step()
    # Calibration/reference set and held-out test subset.
    rng=np.random.default_rng(SEED)
    cal_idx=rng.choice(len(Xtr),size=180,replace=False)
    cal=Xtr[cal_idx]
    test_idx=np.arange(min(360,len(Xte)))
    Xt=Xte[test_idx]; yt=yte[test_idx]
    # Multiple stochastic outputs on calibration and test points.
    A=stochastic_logits(model,cal,24); B=stochastic_logits(model,Xt,80)
    # Robust metric: median pairwise same-seed output distance, with calibration medoid anchor.
    flat=A.mean(0)
    pd=((flat[:,None,:]-flat[None,:,:])**2).sum(2)**.5
    medoid=int(np.argmin(pd.mean(1)))
    dist_anchor=pd[:,medoid]
    # Phi approximation on a log radius grid using empirical ball masses.
    def phi_for(z):
        dd=np.sqrt(((flat-z[None,:])**2).sum(1)); v=np.sqrt(((z-flat[medoid])**2).sum())
        if v < 1e-8: return 0.
        lo=max(v*1e-5,1e-7); rs=np.geomspace(lo,4*v,180)
        mass=np.maximum((dd[:,None] <= rs[None,:]).mean(0),1/len(flat))
        return float(np.trapz(np.log(1/mass),rs))
    phi=np.array([phi_for(z) for z in B.mean(0)])
    va=np.sqrt(((B.mean(0)-flat[medoid])**2).sum(1))
    alpha=1.0; delta=.1
    M0=phi+va*math.log(math.e/delta)
    # Fit a conservative multiplicative C from calibration-like split of test points,
    # using high-sample mean as target stochastic prediction error.
    probs=(np.exp(B-B.max(2,keepdims=True))).sum(2)
    probs= np.exp(B-B.max(2,keepdims=True)); probs/=probs.sum(2,keepdims=True)
    meanp=probs.mean(0); pred=meanp.argmax(1)
    # Actual MC error is deviation of K=80 mean from single stochastic draws.
    point_err=np.sqrt(((probs-meanp[None,:,:])**2).sum(2)).mean(0)
    C=float(np.quantile(point_err/(M0+1e-8),.9))
    M=C*M0
    # Estimate adaptive K from the proposed inverse-square rule, with practical bounds.
    eps=.035; kmin=2; kmax=32
    K=np.clip(np.ceil(2*(M/eps)**2),kmin,kmax).astype(int)
    # Compare mean compute and accuracy for fixed K chosen at the adaptive mean,
    # and a global variance gate with same bounds.
    fixed=int(np.clip(round(K.mean()),kmin,kmax))
    # Evaluate predictions using prefixes of the same stochastic draws.
    def acc_for(ks):
        out=[]
        for i,k in enumerate(ks): out.append(probs[:k,i].mean(0).argmax())
        return float(np.mean(np.array(out)==yt))
    fixed_acc=acc_for(np.full(len(Xt),fixed))
    adaptive_acc=acc_for(K)
    var=np.sqrt(((probs-probs.mean(0)[None,:,:])**2).sum(2).mean(0))
    order=np.argsort(var)
    # Global variance gating uses a single threshold/rank and equal average K.
    globalK=np.full(len(Xt),fixed)
    budget=int(K.sum())
    globalK[:]=kmin
    rem=max(0,budget-kmin*len(Xt)); globalK[order[::-1][:min(len(Xt),rem//(kmax-kmin+1))]] = kmax
    # Fill leftover approximately with intermediate allocation.
    for i in order[::-1]:
        if globalK.sum()>=budget: break
        globalK[i]=min(kmax,globalK[i]+1)
    global_acc=acc_for(globalK)
    rho_phi=float(spearmanr(M,point_err).statistic) if np.std(M)>0 else 0.
    rho_var=float(spearmanr(var,point_err).statistic) if np.std(var)>0 else 0.
    result={"math_sanity":sanity,"n_test":len(Xt),"C":C,"epsilon":eps,
      "adaptive_mean_K":float(K.mean()),"fixed_K":fixed,"global_mean_K":float(globalK.mean()),
      "adaptive_accuracy":adaptive_acc,"fixed_accuracy":fixed_acc,"global_variance_accuracy":global_acc,
      "phi_error_spearman":rho_phi,"variance_error_spearman":rho_var,
      "adaptive_failure_rate_at_80_error":float(np.mean(point_err>M)),
      "M_range":[float(M.min()),float(M.max())]}
    print(json.dumps(result,indent=2))
    with open('results.json','w') as f: json.dump(result,f,indent=2)

if __name__=='__main__': main()