import json, time
import numpy as np
import torch
from torch import nn
from toposcan import grid_tokens, make_transformer

SEED = 19
np.random.seed(SEED); torch.manual_seed(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if DEVICE == "cuda": torch.cuda.get_device_name(0)
except Exception:
    DEVICE = "cpu"


def cycle_graph(lengths):
    edges=[]; off=0
    for L in lengths:
        for i in range(L): edges.append((off+i, off+(i+1)%L))
        off += L
    return off, edges


def make_data(n_each=160):
    # Every graph has 12 vertices and every vertex degree 2.
    # Label 0: one 12-cycle (beta1=1); label 1: two 6-cycles (beta1=2).
    graphs=[]
    for y, lengths in [(0,[12]), (1,[6,6])]:
        for _ in range(n_each):
            n,e=cycle_graph(lengths)
            # identical filtration keeps topology active in one shared window
            h=np.full(n, .5, dtype=np.float32)
            graphs.append((n,e,h,y))
    rng=np.random.default_rng(SEED); rng.shuffle(graphs)
    return graphs


def features(graphs):
    degree=[]; topo=[]; labels=[]; prep=0.
    grid=np.linspace(0.,1.,9)
    for n,e,h,y in graphs:
        d=np.zeros(n, dtype=np.float32)
        for u,v in e: d[u]+=1; d[v]+=1
        degree.append(np.bincount(d.astype(int), minlength=4).astype(np.float32)/n)
        t0=time.perf_counter()
        tok=grid_tokens(n,e,h,grid,m=4,stride=2)
        prep += time.perf_counter()-t0
        topo.append(tok); labels.append(y)
    return np.asarray(degree), topo, np.asarray(labels), prep


def train_mlp(xtr,ytr,xte,yte, epochs=100):
    model=nn.Sequential(nn.Linear(xtr.shape[1],16),nn.ReLU(),nn.Linear(16,2)).to(DEVICE)
    opt=torch.optim.Adam(model.parameters(),lr=.03)
    X=torch.tensor(xtr,dtype=torch.float32,device=DEVICE); Y=torch.tensor(ytr,device=DEVICE)
    for _ in range(epochs):
        opt.zero_grad(); loss=nn.functional.cross_entropy(model(X),Y); loss.backward(); opt.step()
    with torch.no_grad():
        pred=model(torch.tensor(xte,dtype=torch.float32,device=DEVICE)).argmax(1).cpu().numpy()
    return float(np.mean(pred==yte)), sum(p.numel() for p in model.parameters())


def train_topo(ttr,ytr,tte,yte,epochs=100):
    # Batch padding is unnecessary here: all token sequences have the same length.
    model=nn.Sequential(make_transformer(width=16,heads=4,layers=1),nn.Linear(16,2)).to(DEVICE)
    opt=torch.optim.Adam(model.parameters(),lr=.02)
    X=torch.tensor(np.asarray(ttr),dtype=torch.float32,device=DEVICE); Y=torch.tensor(ytr,device=DEVICE)
    for _ in range(epochs):
        opt.zero_grad(); loss=nn.functional.cross_entropy(model(X),Y); loss.backward(); opt.step()
    with torch.no_grad():
        pred=model(torch.tensor(np.asarray(tte),dtype=torch.float32,device=DEVICE)).argmax(1).cpu().numpy()
    return float(np.mean(pred==yte)), sum(p.numel() for p in model.parameters())


def main():
    data=make_data(); cut=round(.7*len(data))
    tr,te=data[:cut],data[cut:]
    dtr,ttr,ytr,p1=features(tr); dte,tte,yte,p2=features(te)
    # perturbation test: filtration shifts do not change the relative topology here;
    # use a separate score perturbation and recompute tokens.
    acc_d,par_d=train_mlp(dtr,ytr,dte,yte)
    acc_t,par_t=train_topo(ttr,ytr,tte,yte)
    # A cheap direct topological classifier isolates representation power.
    ztr=np.asarray([a[:,1].max() for a in ttr])[:,None]
    zte=np.asarray([a[:,1].max() for a in tte])[:,None]
    acc_direct,par_direct=train_mlp(ztr,ytr,zte,yte)
    results={"device":DEVICE,"graphs":len(data),"degree_hist_unique":int(np.unique(np.concatenate([dtr,dte]),axis=0).shape[0]),
             "baseline_degree_accuracy":acc_d,"toposcan_transformer_accuracy":acc_t,
             "direct_beta1_accuracy":acc_direct,"baseline_params":par_d,"toposcan_params":par_t,
             "preprocess_seconds_total":p1+p2,"preprocess_seconds_per_graph":(p1+p2)/len(data),
             "note":"All graphs have identical degree histogram; topology labels are beta1=1 vs 2."}
    print(json.dumps(results,indent=2))

if __name__=='__main__': main()