import numpy as np

META = {"name": "relational_block_graph", "domain": "graph-nn", "description": "Multi-relational stochastic-block edge prediction with node features; targets are sampled held-out relation edges."}

def get_dataset(seed, n_train=400, n_test=120):
    rng = np.random.default_rng(seed)
    n_nodes, d, m, r = 36, 8, 3, 2
    z = rng.integers(m, size=n_nodes)
    centers = rng.normal(0, 1, (m, d))
    X = centers[z] + 0.35 * rng.normal(size=(n_nodes, d))
    A = np.array([[[.78,.16,.32],[.16,.72,.20],[.32,.20,.70]],
                  [[.65,.25,.15],[.25,.72,.28],[.15,.28,.62]]], dtype=np.float32)
    pairs, ys = [], []
    for _ in range(n_train + n_test):
        u, v = rng.integers(n_nodes, size=2)
        k = int(rng.integers(r))
        y = float(rng.random() < A[k, z[u], z[v]])
        pairs.append(np.concatenate([X[u], X[v], [k]]))
        ys.append(y)
    pairs = np.asarray(pairs, dtype=np.float32)
    ys = np.asarray(ys, dtype=np.float32)[:, None]
    mu = pairs[:n_train, :2*d].mean(0)
    sd = pairs[:n_train, :2*d].std(0) + 1e-6
    pairs[:, :2*d] = (pairs[:, :2*d] - mu) / sd
    return {"xtr": pairs[:n_train], "ytr": ys[:n_train],
            "xte": pairs[n_train:], "yte": ys[n_train:],
            "task": "regression", "metric": "mse", "input_shape": (2*d+1,),
            "out_dim": 1}