import sys, json
import numpy as np
import torch
import torch.nn as nn
ROOT='/home/maxwelhelp/all/math2nn'
if ROOT not in sys.path: sys.path.insert(0,ROOT)
from bench import train_model, make_report, sweep_baseline
from bench.protocol import DEFAULT_SEEDS, SWEEP_SEEDS
META={'name':'wedge_graph_regression','domain':'graph-nn','description':'Graph-level regression from node signals on connected cycle-plus-chord graphs; adaptive shortest-path wedge pooling preserves connected regions.'}
N,D,M=18,4,6
ADJ=[set() for _ in range(N)]
for i in range(N):
    ADJ[i].update(((i-1)%N,(i+1)%N))
for a,b in [(0,9),(3,12),(6,15)]: ADJ[a].add(b); ADJ[b].add(a)
ADJ=[sorted(s) for s in ADJ]

def get_dataset(seed,n_train=400,n_test=200):
    rng=np.random.default_rng(seed); n=n_train+n_test; phase=rng.uniform(0,2*np.pi,n); amp=rng.uniform(.5,1.5,n); x=np.zeros((n,N,D),np.float32); idx=np.arange(N)
    for k in range(n):
        x[k,:,0]=amp[k]*np.sin(2*np.pi*idx/N+phase[k])+.08*rng.normal(size=N)
        x[k,:,1]=np.where(idx<9,rng.uniform(-1,1),rng.uniform(-1,1))+.08*rng.normal(size=N)
        x[k,:,2]=np.cos(4*np.pi*idx/N+phase[k])+.08*rng.normal(size=N); x[k,:,3]=rng.normal(size=N)
    edges=[(a,b) for a in range(N) for b in ADJ[a] if b>a]; e=np.zeros(n)
    for a,b in edges: e+=(x[:,a,0]-x[:,b,0])**2+.5*(x[:,a,1]-x[:,b,1])**2
    y=(e/len(edges)+.25*x[:,:,0].mean(1)+.1*x[:,:,2].mean(1)).astype(np.float32)[:,None]
    return {'xtr':torch.from_numpy(x[:n_train]),'ytr':torch.from_numpy(y[:n_train]),'xte':torch.from_numpy(x[n_train:]),'yte':torch.from_numpy(y[n_train:]),'task':'regression','metric':'mse','input_shape':(M,D),'out_dim':1}

def bfs(region):
    R=set(region); allD={}
    for s in R:
        ds={s:0}; q=[s]
        for v in q:
            for z in ADJ[v]:
                if z in R and z not in ds: ds[z]=ds[v]+1; q.append(z)
        allD[s]=ds
    return allD

def wedge_regions(X):
    regions=[list(range(N))]
    while len(regions)<M:
        best=None
        for ri,R in enumerate(regions):
            if len(R)<2: continue
            ds=bfs(R); mu=X[R].mean(0); base=((X[R]-mu)**2).sum()
            for u in R:
                for w in R:
                    if u>=w: continue
                    A=[v for v in R if ds[u][v]<=ds[w][v]]; B=[v for v in R if ds[u][v]>ds[w][v]]
                    if not A or not B: continue
                    gain=base-((X[A]-X[A].mean(0))**2).sum()-((X[B]-X[B].mean(0))**2).sum()
                    if best is None or gain>best[0]: best=(gain,ri,A,B)
        if best is None or best[0]<=1e-10: break
        _,ri,A,B=best; regions.pop(ri); regions.extend((A,B))
    return sorted(regions,key=min)

def pool_array(x,mode,seed):
    out=np.empty((len(x),M,D),np.float32)
    if mode=='random':
        rng=np.random.default_rng(seed); labels=np.repeat(np.arange(M),int(np.ceil(N/M)))[:N]; rng.shuffle(labels); regs=[np.flatnonzero(labels==i) for i in range(M)]
        for k in range(len(x)):
            for i,r in enumerate(regs): out[k,i]=x[k,r].mean(0)
    else:
        for k in range(len(x)):
            regs=wedge_regions(x[k])
            for i,r in enumerate(regs): out[k,i]=x[k,r].mean(0)
            if len(regs)<M: out[k,len(regs):]=out[k,len(regs)-1]
    return torch.from_numpy(out)

class PoolNet(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(M*D,64),nn.ReLU(),nn.Linear(64,1))
    def forward(self,x): return self.net(x.reshape(x.shape[0],-1))

def train_one(mode,lr,seed,epochs=18):
    torch.manual_seed(seed); np.random.seed(seed); ds=get_dataset(seed,400,200)
    ds['xtr']=pool_array(ds['xtr'].numpy(),mode,seed); ds['xte']=pool_array(ds['xte'].numpy(),mode,seed+10000)
    net,metric,_=train_model(PoolNet(),ds,epochs=epochs,lr=lr,batch=128,log=lambda *a:None)
    net=net.cpu()
    with torch.no_grad(): pred=net(ds['xte']).numpy().ravel(); obs=ds['yte'].numpy().ravel()
    return float(metric),{'pred_std':float(pred.std()),'obs_std':float(obs.std()),'pred_obs_corr':float(np.corrcoef(pred,obs)[0,1])}

def eval_cfg(mode,cfg,seeds=DEFAULT_SEEDS):
    scores=[]; sig=[]
    for s in seeds:
        v,z=train_one(mode,float(cfg['lr']),int(s),int(cfg.get('epochs',18))); scores.append(v); sig.append(z)
    return {'mean':float(np.mean(scores)),'std':float(np.std(scores)),'per_seed':scores,'n':len(scores),'signature':sig}

def main():
    grid=[{'lr':1e-3,'epochs':18},{'lr':3e-3,'epochs':18},{'lr':1e-2,'epochs':18}]
    base=sweep_baseline(lambda c: (lambda s: train_one('random',c['lr'],s,c['epochs'])[0]),grid,seeds=SWEEP_SEEDS)
    # Evaluate the same union of configs for the idea, then select its best.
    idea_grid={str(c['lr']):eval_cfg('wedge',c) for c in grid}
    best=min(idea_grid,key=lambda k:idea_grid[k]['mean']); idea=idea_grid[best]
    # Mechanism signature is measured from trained models, not an identity.
    sig={'predicted':'graph-aware connected pooling should retain graph-smooth signal (higher prediction/observation correlation than random pooling)','observed':{'baseline_best_corr_mean':float(np.mean([train_one('random',base['best_cfg']['lr'],s,base['best_cfg']['epochs'])[1]['pred_obs_corr'] for s in DEFAULT_SEEDS])),'idea_corr_mean':float(np.mean([z['pred_obs_corr'] for z in idea['signature']]))},'confirmed':False}
    report=make_report('wedge_graph_regression','mlp_tiny',base,idea,{'mechanism_signature':sig,'custom_track':{'name':META['name'],'file':'graph_wedge_bench.py','domain':META['domain']},'idea_grid':idea_grid})
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()