import numpy as np
META = {'name':'corrupted_root_graph','domain':'graph-nn','description':'Root localization on permuted recursive trees with shortcut corruption.'}
def _one(rng,n=32,shortcuts=8):
    edges=[(v,int(rng.integers(v))) for v in range(1,n)]
    used={tuple(sorted(e)) for e in edges}
    for _ in range(shortcuts):
        for _ in range(100):
            a,b=map(int,rng.integers(n,size=2)); e=tuple(sorted((a,b)))
            if a!=b and e not in used: break
        used.add(e); edges.append(e)
    perm=rng.permutation(n); A=np.zeros((n,n),np.float32)
    for a,b in edges: a,b=int(perm[a]),int(perm[b]); A[a,b]=A[b,a]=1
    return A,int(np.where(perm==0)[0][0])
def get_dataset(seed,n_train,n_test):
    rng=np.random.default_rng(seed); n=32
    def make(k):
        x=np.empty((k,n,n),np.float32); y=np.empty(k,np.int64)
        for i in range(k): x[i],y[i]=_one(rng,n)
        return x,y
    xtr,ytr=make(n_train); xte,yte=make(n_test)
    return {'xtr':xtr,'ytr':ytr,'xte':xte,'yte':yte,'task':'classification','metric':'cross_entropy'}