import numpy as np

META = {"name": "torus_cayley_graph", "domain": "graph-nn", "description": "Nodewise regression on labeled Z_6 x Z_6 Cayley torus graphs; targets are a fixed local relation convolution."}
N = 6
EDGES, RELS = [], []
for x in range(N):
    for y in range(N):
        u = x * N + y
        for r, (dx, dy) in enumerate(((1, 0), (-1, 0), (0, 1), (0, -1))):
            v = ((x + dx) % N) * N + ((y + dy) % N)
            EDGES.append((u, v)); RELS.append(r)
EDGES = np.asarray(EDGES, dtype=np.int64)
RELS = np.asarray(RELS, dtype=np.int64)
COEF = np.asarray([1.0, -0.7, 0.45, -1.15], dtype=np.float32)

def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    def make(n):
        x = rng.normal(size=(n, N*N, 4)).astype(np.float32)
        y = np.zeros((n, N*N, 1), dtype=np.float32)
        for (u, v), r in zip(EDGES, RELS):
            y[:, v, 0] += COEF[r] * x[:, u, 0]
        y += rng.normal(0, 0.03, y.shape).astype(np.float32)
        return x, y
    xtr, ytr = make(n_train); xte, yte = make(n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}