import sys, json, random, math, importlib.util
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
import bench
from bench import evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))


def register_local_track():
    p = Path(__file__).with_name('graph_track.py')
    spec = importlib.util.spec_from_file_location('local_graph_track', p)
    mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod)
    # Do not edit bench: register only in this process for the required smoke test.
    bench.data._CUSTOM_CACHE = {mod.META['name']: mod}
    d = bench.get_dataset(mod.META['name'], seed=0, n_train=16, n_test=8)
    assert d['task'] == 'classification' and d['xtr'].shape[0] == 16
    return mod.META['name']


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)


def make_graph(seed):
    rng=np.random.default_rng(seed); n=240; k=120
    y=np.r_[np.zeros(k,dtype=np.int64),np.ones(k,dtype=np.int64)]
    same=y[:,None]==y[None,:]
    tri=np.triu(rng.random((n,n)) < np.where(same,.18,.012),1)
    A=np.zeros((n,n),np.float32); A[tri]=1.; A[tri.T]=1.
    x=rng.normal(0,.7,(n,8)).astype(np.float32); s=2*y.astype(np.float32)-1
    x[:,0]=.18*s; x[:,1]=.12*s
    flip=rng.random(n)<.12; y=y.copy(); y[flip]=1-y[flip]
    perm=rng.permutation(n); return A,x,y,perm[:160],perm[160:]


def lap_columns(A):
    d=A.sum(1); cols=[]; norms=[]
    for i in range(len(A)):
        z={i:float(d[i])}
        for j in np.flatnonzero(A[i]): z[int(j)]=z.get(int(j),0.)-float(A[i,j])
        cols.append(z); norms.append(float(d[i]**2+(A[i]**2).sum()))
    return cols,np.asarray(norms)


def coh(i,j,cols,norms):
    if len(cols[i])>len(cols[j]): i,j=j,i
    dot=sum(v*cols[j].get(q,0.) for q,v in cols[i].items())
    return abs(dot)/(math.sqrt(norms[i]*norms[j])+1e-12)


def coherent(candidates,batch,cols,norms,rng):
    pool=list(candidates); first=pool[rng.randrange(len(pool))]; out=[first]
    rem=set(pool); rem.remove(first)
    while len(out)<batch:
        q=min(rem,key=lambda i:max(coh(i,j,cols,norms) for j in out))
        out.append(q); rem.remove(q)
    return np.asarray(out,dtype=np.int64)


class GCN(nn.Module):
    def __init__(self):
        super().__init__(); self.l1=nn.Linear(8,32); self.l2=nn.Linear(32,2)
    def forward(self,x,P): return self.l2(torch.relu(self.l1(P@x)))


def train_one(seed,cfg,idea,capture=False):
    seed_all(seed); A,x,y,tr,te=make_graph(seed); n=len(y)
    P=torch.tensor(A/np.maximum(A.sum(1)[:,None],1)+np.eye(n,dtype=np.float32))
    xt,yt=torch.tensor(x),torch.tensor(y); net=GCN()
    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['wd'])
    cols,norms=lap_columns(A); rng=random.Random(seed+991)
    counts=np.zeros(n)
    if idea:
        for _ in range(40):
            cc=rng.sample(list(tr),cfg['cand']); counts[coherent(cc,cfg['batch'],cols,norms,rng)]+=1
        p=np.maximum(counts/40.,1/40.)
    else: p=np.full(n,cfg['batch']/len(tr))
    for _ in range(cfg['epochs']):
        if idea: sel=coherent(rng.sample(list(tr),cfg['cand']),cfg['batch'],cols,norms,rng)
        else: sel=np.asarray(rng.sample(list(tr),cfg['batch']))
        pred=net(xt,P); losses=nn.functional.cross_entropy(pred[sel],yt[sel],reduction='none')
        loss=(losses*torch.tensor(1/(n*p[sel]),dtype=torch.float32)).sum() if idea else losses.mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad(): err=float((net(xt,P)[te].argmax(1)!=yt[te]).float().mean())
    if not capture:return err
    cov=[]; uc=[]
    for _ in range(200):
        cov.append(len(set((coherent(rng.sample(list(tr),cfg['cand']),cfg['batch'],cols,norms,rng)>=120).astype(int))))
        ss=rng.sample(list(tr),cfg['batch']); uc.append(len(set((np.asarray(ss)>=120).astype(int))))
    predicted=2*(1-math.comb(80,cfg['cand'])/math.comb(160,cfg['cand'])) if cfg['cand']<=80 else 2.
    return {'metric':err,'observed_coherent_coverage':float(np.mean(cov)),
            'observed_uniform_coverage':float(np.mean(uc)),
            'predicted_candidate_pool_coverage':float(predicted),
            'confirmed':bool(abs(np.mean(cov)-predicted)<.12)}


def main():
    track=register_local_track()
    lrs=[.001,.003,.006]
    grid=[{'lr':lr,'wd':wd,'epochs':18,'batch':32,'cand':64} for lr in lrs for wd in [0.,1e-4]]
    base=sweep_baseline(lambda c:lambda s:train_one(s,c,False),grid,seeds=(0,1,2,3))
    runs=[(dict(base['best_cfg'],lr=lr),evaluate(lambda s,c=dict(base['best_cfg'],lr=lr):train_one(s,c,True),seeds=SEEDS)) for lr in lrs]
    cfg,idea=min(runs,key=lambda z:z[1]['mean'])
    sig=train_one(0,cfg,True,True)
    report=make_report(track,'local_gcn',base,idea,extra={'custom_track':{'name':track,'file':'graph_track.py','domain':'graph-nn'},'observed_best_cfg':cfg,'idea_sweep':[{'cfg':c,'result':r} for c,r in runs],'mechanism_signature':sig})
    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))

if __name__=='__main__': main()