"""Small deterministic attributed relational-graph benchmark for stage-2."""
import numpy as np

META = {
    "name": "relational_graph_classification",
    "domain": "graph-nn",
    "description": "Permutation-augmented attributed graph classification with dense relational kernels.",
}


def _make(rng, n):
    out, ys = [], []
    for _ in range(int(n)):
        y = int(rng.randint(0, 2))
        p = 0.22 if y == 0 else 0.58
        a = (rng.rand(8, 8) < p).astype(np.float32)
        a = np.triu(a, 1)
        a = a + a.T
        np.fill_diagonal(a, 0.0)
        node = rng.normal(0.0, 0.12, size=(8, 1)).astype(np.float32)
        node += 0.08 if y else -0.08
        out.append(np.concatenate([a, node], axis=1))
        ys.append(y)
    return np.asarray(out, dtype=np.float32), np.asarray(ys, dtype=np.int64)


def get_dataset(seed, n_train, n_test):
    xtr, ytr = _make(np.random.RandomState(int(seed)), n_train)
    xte, yte = _make(np.random.RandomState(int(seed) + 100003), n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "classification", "metric": "err", "out_dim": 2}