import json, math, random
from pathlib import Path
import numpy as np
from sklearn.neighbors import NearestNeighbors
from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn

SEED = 523
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def graph_from_x(x, k=10, sigma=1.0):
    n = len(x)
    d = ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
    nnidx = np.argsort(d, axis=1)[:, 1:k+1]
    W = np.zeros((n, n), dtype=np.float64)
    for i in range(n):
        for j in nnidx[i]: W[i, j] = math.exp(-d[i, j] / (2*sigma*sigma))
    W = np.maximum(W, W.T); np.fill_diagonal(W, 0)
    return W

def cut_ratio(W, y):
    return float((W * (y[:,None] != y[None,:])).sum() / (W.sum() + 1e-12))

def weights_from_cuts(cuts, beta=12.0, qmin=0.0):
    z = -beta*np.asarray(cuts); z -= z.max(); q = np.exp(z); q /= q.sum()
    if qmin:
        q = np.maximum(q, qmin); q /= q.sum()
    return q

def verify_math():
    # Two clearly separated groups and a graph with only within-group edges.
    y = np.array([0,0,0,1,1,1])
    Wgood = np.zeros((6,6)); Wbad = np.ones((6,6))-np.eye(6)
    for a in range(3):
        for b in range(3): Wgood[a,b+3] = Wgood[b+3,a] = 0
    for i in range(3):
        for j in range(3):
            if i != j: Wgood[i,j] = Wgood[i+3,j+3] = 1
    rg, rb = cut_ratio(Wgood,y), cut_ratio(Wbad,y)
    q = weights_from_cuts([rg, rb], beta=10)
    assert rg < rb and q[0] > q[1] and abs(q.sum()-1) < 1e-12
    # Directly check the stated ratio is invariant to global graph scaling.
    assert abs(cut_ratio(7.3*Wbad,y)-rb) < 1e-12
    return {"good_cut":rg, "bad_cut":rb, "q_beta10":q.tolist(), "scale_invariance_error":abs(cut_ratio(7.3*Wbad,y)-rb)}

class Net(nn.Module):
    def __init__(self):
        super().__init__(); self.f = nn.Sequential(nn.Linear(2,32),nn.ReLU(),nn.Linear(32,32),nn.ReLU(),nn.Linear(32,3))
    def forward(self,x): return self.f(x)

def make_data(seed=SEED):
    rng=np.random.RandomState(seed)
    centers=np.array([[-2.,0.],[2.,0.],[0.,2.8]])
    xs=[]; ys=[]
    for c in range(3):
        xs.append(centers[c]+rng.randn(90,2)*0.62); ys += [c]*90
    return np.vstack(xs).astype('float32'), np.array(ys,dtype=np.int64)

def train_once(mode, X, y, graphs, cuts, seed):
    torch.manual_seed(seed)
    # fixed 12 labeled points/class, remaining nodes participate only in graph term
    labeled=np.concatenate([np.where(y==c)[0][:12] for c in range(3)])
    xt=torch.tensor(X); yt=torch.tensor(y); lt=torch.tensor(labeled)
    model=Net(); opt=torch.optim.Adam(model.parameters(),lr=0.025)
    if mode=='uniform': q=np.ones(len(graphs))/len(graphs)
    elif mode=='cut': q=weights_from_cuts(cuts,beta=12.)
    else: raise ValueError(mode)
    W=sum(float(qi)*Wi for qi,Wi in zip(q,graphs))
    # normalize graph penalty to comparable scale across policies
    Wt=torch.tensor(W,dtype=torch.float32); denom=Wt.sum().clamp_min(1e-8)
    for _ in range(260):
        opt.zero_grad(); logits=model(xt); p=torch.softmax(logits,1)
        sup=nn.functional.cross_entropy(logits[lt],yt[lt])
        dif=(p[:,None,:]-p[None,:,:]).pow(2).sum(-1)
        reg=(Wt*dif).sum()/denom
        (sup+0.8*reg).backward(); opt.step()
    with torch.no_grad():
        pred=model(xt).argmax(1).numpy(); p=torch.softmax(model(xt),1).numpy()
    return {"accuracy":float((pred==y).mean()),"labeled_accuracy":float((pred[labeled]==y[labeled]).mean()),"graph_reg":float(reg.item()),"q":q.tolist(),"pred_cut_uniform_graph":cut_ratio(W,pred)}

def main():
    mathcheck=verify_math(); X,y=make_data()
    rng=np.random.RandomState(SEED)
    # Identity and mild noise preserve neighborhoods; shuffled coordinates intentionally mix labels.
    policies=[X, X+rng.randn(*X.shape).astype('float32')*0.18, X[rng.permutation(len(X))]]
    graphs=[graph_from_x(z,k=10,sigma=0.8) for z in policies]
    cuts=[cut_ratio(W,y) for W in graphs]
    results={m:train_once(m,X,y,graphs,cuts,SEED+17) for m in ['uniform','cut']}
    out={"math_check":mathcheck,"policy_cuts":cuts,"cut_weights":weights_from_cuts(cuts,12).tolist(),"results":results}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()