import numpy as np

META = {
    'name': 'soft_graph_connectivity',
    'domain': 'graph-nn',
    'description': 'Synthetic soft graph edge-weight regression with eight-node adjacency structure.'
}

def get_dataset(seed, n_train, n_test):
    def make(rng, n):
        x = np.empty((n, 8, 2), dtype=np.float32)
        y = np.empty((n, 1), dtype=np.float32)
        ii, jj = np.triu_indices(8, 1)
        for b in range(n):
            ang = np.linspace(0, 2*np.pi, 8, endpoint=False) + rng.normal(0, .06, 8)
            rad = 1.0 + rng.normal(0, .05, 8)
            pts = np.stack((rad*np.cos(ang), rad*np.sin(ang)), axis=1)
            pts += rng.normal(0, .035, pts.shape)
            w = np.zeros((8, 8), dtype=np.float32)
            for i in range(8):
                for j in range(i):
                    d = np.linalg.norm(pts[i] - pts[j])
                    ring = min((i-j) % 8, (j-i) % 8) == 1
                    value = np.exp(-d*d/1.8) * (1.0 if ring else .42)
                    w[i, j] = w[j, i] = max(.01, value + rng.normal(0, .012))
            x[b] = pts
            y[b, 0] = w[ii, jj].mean()
        return x, y
    xtr, ytr = make(np.random.RandomState(seed), n_train)
    xte, yte = make(np.random.RandomState(seed + 5000), n_test)
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'task': 'regression', 'metric': 'mse', 'out_dim': 1}