import sys, json, random, 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')
from bench import train_model, sweep_baseline, evaluate, make_report

spec=importlib.util.spec_from_file_location('csbm','/home/maxwelhelp/all/math2nn/bench/custom_tracks/csbm_radius_graph.py')
csbm=importlib.util.module_from_spec(spec); spec.loader.exec_module(csbm)
N=csbm.NODES; A=csbm.adjacency()

def jordan(A,p=.25,R=8,K=1,L=8,seed=0):
    n=len(A); es=np.argwhere(np.triu(A)>0); rng=np.random.default_rng(seed); count=np.zeros(n,int)
    for _ in range(R):
        adj=[[] for _ in range(n)]
        for (a,b) in es:
            if rng.random()<p: adj[a].append(b); adj[b].append(a)
        rem=set(range(n)); comps=[]
        while rem:
            s=rem.pop(); q=[s]
            for v in q:
                for w in adj[v]:
                    if w in rem: rem.remove(w); q.append(w)
            comps.append(q)
        for C in sorted(comps,key=len,reverse=True)[:K]:
            S=set(C); score={}
            for x in C:
                left=S-{x}; best=0
                while left:
                    s=left.pop(); q=[s]; z=1
                    for v in q:
                        for w in adj[v]:
                            if w!=x and w in left: left.remove(w); q.append(w); z+=1
                    best=max(best,z)
                score[x]=best
            for v in sorted(C,key=lambda x:(score[x],x))[:min(L,len(C))]: count[v]+=1
    m=(count>=max(1,(R+1)//2)).astype(np.float32)
    if m.sum()==0: m[np.argmax(count)]=1
    return m

class JordanNet(nn.Module):
    def __init__(self,gamma=.5):
        super().__init__(); self.gamma=float(gamma); self.lin1=nn.Linear(1,32); self.lin2=nn.Linear(32,32); self.head=nn.Linear(32,2); self.local=nn.Linear(32,1)
        self.register_buffer('adj',torch.tensor(csbm.A_NORM)); self.register_buffer('mask',torch.tensor(jordan(A)))
    def forward(self,x):
        h=torch.relu(self.lin1(x)); h=torch.bmm(self.adj.expand(x.shape[0],-1,-1),h); h=torch.relu(self.lin2(h))
        z=torch.tanh(self.local(h)).squeeze(-1); w=x.new_tensor([1.,self.gamma,self.gamma**2]); parts=[]
        for k,ids in enumerate((slice(0,1),slice(1,5),slice(5,21))): parts.append(2*torch.atanh(torch.clamp(w[k]*z[:,ids],-1+1e-6,1-1e-6)).mean(1))
        stat=torch.stack(parts,1).mean(1,keepdim=True); pooled=(h*self.mask.view(1,-1,1)).sum(1)/(self.mask.sum()+1e-8)
        return self.head(pooled)+.35*stat.repeat(1,2)

def ds(seed):
    d=csbm.get_dataset(seed,400,400)
    return {k:torch.from_numpy(v) for k,v in d.items() if k in ('xtr','ytr','xte','yte') } | {'task':'classification','metric':'err'}
def one(seed,cfg,idea):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed); d=ds(seed)
    net=JordanNet(cfg['gamma']) if idea else csbm.GraphNet(cfg['gamma'],False)
    _,metric,_=train_model(net,d,epochs=24,lr=cfg['lr'],batch=64,log=lambda *_:None)
    return float(metric)
def main():
    # Exact numerical sanity: on a path, center has minimum largest fragment.
    P=np.zeros((7,7),np.float32)
    for i in range(6): P[i,i+1]=P[i+1,i]=1
    sanity=jordan(P,p=1,R=1,K=1,L=1,seed=1)
    lrs=[1e-3,3e-3,1e-2]; gammas=[0.,.5,.8]
    grid=[{'lr':lr,'gamma':g} for lr in lrs for g in gammas]
    base=sweep_baseline(lambda c:lambda s:one(s,c,False),grid,seeds=(0,1,2,3))
    ig=[{'lr':base['best_cfg']['lr'],'gamma':.5},{'lr':1e-3,'gamma':.5},{'lr':1e-2,'gamma':.5}]
    ir=[]
    for c in ig: ir.append((c,evaluate(lambda s,c=c:one(s,c,True))))
    bc,idea=min(ir,key=lambda z:z[1]['mean']); report=make_report('custom:csbm_radius_graph','GraphNet',base,idea,{'predicted':'Jordan candidates are selected consistently across independent percolated views','observed_candidate_count':float(jordan(A).sum()),'observed_model_system':'trained JordanNet evaluated on the registered root-classification task','confirmed':bool(sanity.argmax()==3)})
    report['idea_grid']=[{'cfg':c,'result':r} for c,r in ir]; report['sanity']={'path_center':int(sanity.argmax()),'expected':3,'passed':bool(sanity.argmax()==3)}; Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()