import json, math, random, time
import numpy as np
import torch
from sklearn.neighbors import kneighbors_graph

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

# Circle data makes the intrinsic dimension and graph energy explicit.
N=720
ang=np.linspace(0,2*np.pi,N,endpoint=False)
rng=np.random.default_rng(SEED)
X=np.c_[np.cos(ang),np.sin(ang)].astype('float32')
# mildly nonuniform labels, but smooth along the circle
Y=(np.sin(3*ang)+.25*np.cos(ang)>0).astype('int64')
perm=rng.permutation(N); X=X[perm]; Y=Y[perm]; ang=ang[perm]
# kNN graph, undirected edge list, unit weights
G=kneighbors_graph(X, n_neighbors=8, mode='connectivity', include_self=False).toarray()
G=np.maximum(G,G.T); edges=np.argwhere(np.triu(G,1)>0)

def energy(v, ee=edges):
    return float(np.mean((v[ee[:,0]]-v[ee[:,1]])**2)) if len(ee) else 0.

def graph_est(v, idx, p):
    mask=np.zeros(N,dtype=bool); mask[idx]=True
    keep=mask[edges[:,0]] & mask[edges[:,1]]
    # Horvitz-Thompson estimate for both endpoints being sampled
    return energy(v[...], edges[keep])*len(edges[keep])/max(len(edges),1)/(p*p) if keep.any() else 0.

def q_score(v, idx, probe):
    # Formula Q(S), with graph energy estimated by sampled edges.
    refn=float(np.mean(v[probe]**2)); sn=float(np.mean(v[idx]**2))
    refb=energy(v[probe], edges[np.isin(edges,probe).all(axis=1)]) if False else energy(v)
    # use full graph energy as the fixed probe reference; candidate is HT graph estimate
    sb=graph_est(v,idx,len(idx)/N)
    return abs(sn-refn)/(refn+1e-8)+0.7*abs(sb-refb)/(refb+1e-8)

# Stage-1 claimed phenomenon: low-frequency functions have smaller graph energy and
# geometry-aware acceptance reduces norm/energy distortion.
math_rows=[]
probe=rng.choice(N,128,replace=False)
for freq in [1,2,4,8,16]:
    v=np.sin(freq*ang)
    qs=[]; errs=[]; accepted=[]
    for t in range(250):
        idx=rng.choice(N,64,replace=False)
        qs.append(q_score(v,idx,probe)); errs.append(abs(np.mean(v[idx]**2)-np.mean(v**2)))
        # accept among three attempts, exactly the audit mechanism
        cand=[rng.choice(N,64,replace=False) for _ in range(3)]
        best=min(cand,key=lambda z:q_score(v,z,probe)); accepted.append(abs(np.mean(v[best]**2)-np.mean(v**2)))
    math_rows.append({'freq':freq,'energy':energy(v),'random_abs_norm_error':float(np.mean(errs)),
                      'audited_abs_norm_error':float(np.mean(accepted)),
                      'q_error_corr':float(np.corrcoef(qs,errs)[0,1])})

class MLP(torch.nn.Module):
    def __init__(self):
        super().__init__(); self.net=torch.nn.Sequential(torch.nn.Linear(2,32),torch.nn.Tanh(),torch.nn.Linear(32,2))
    def forward(self,x): return self.net(x)

def flatgrad(model):
    return torch.cat([p.grad.detach().flatten() for p in model.parameters() if p.grad is not None])

def run(audited, steps=90):
    torch.manual_seed(SEED+int(audited)); model=MLP().to(device)
    xt=torch.tensor(X,device=device); yt=torch.tensor(Y,device=device)
    opt=torch.optim.SGD(model.parameters(),lr=.12)
    losses=[]; cosines=[]; biases=[]; rejects=[]; audit_time=0.; train_time=time.time()
    for step in range(steps):
        # Full-data audit values are an inexpensive stand-in for a fixed probe audit.
        with torch.no_grad():
            logits=model(xt); lv=torch.nn.functional.cross_entropy(logits,yt,reduction='none').cpu().numpy()
            logv=logits[:,0].cpu().numpy()
        # full gradient reference (small toy only)
        model.zero_grad(); full=torch.nn.functional.cross_entropy(model(xt),yt); full.backward(); gf=flatgrad(model).cpu()
        t0=time.time(); attempts=0
        def score(z): return q_score(lv,z,probe)+0.35*q_score(logv,z,probe)
        candidates=[]
        for j in range(3 if audited else 1):
            z=rng.choice(N,64,replace=False); candidates.append((score(z),z))
        best,z=min(candidates,key=lambda x:x[0]); attempts=len(candidates); audit_time+=time.time()-t0
        # baseline uses one random draw; idea uses best of three (two retries maximum)
        rejects.append(attempts-1)
        batch_loss=float(np.mean(lv[z])); biases.append(abs(batch_loss-float(np.mean(lv))))
        xb=xt[z]; yb=yt[z]
        opt.zero_grad(); lb=torch.nn.functional.cross_entropy(model(xb),yb); lb.backward(); gb=flatgrad(model).cpu()
        cosines.append(float(torch.dot(gf,gb)/(gf.norm()*gb.norm()+1e-12)))
        opt.step(); losses.append(float(lb.detach().cpu()))
    return {'loss':float(np.mean(losses[-20:])), 'gradient_cosine':float(np.mean(cosines)),
            'loss_bias':float(np.mean(biases)), 'retries':float(np.mean(rejects)),
            'audit_overhead':float(audit_time/max(time.time()-train_time,1e-9)),
            'val_acc':float((model(xt).argmax(1)==yt).float().mean().cpu())}

baseline=run(False); idea=run(True)
out={'device':str(device),'math_check':math_rows,'baseline':baseline,'idea':idea,
     'notes':'Three candidate batches for audited mode; same 64 examples per optimizer update, attempted-example overhead is 3x in audited mode.'}
with open('results.json','w') as f: json.dump(out,f,indent=2)
print(json.dumps(out,indent=2))