import numpy as np

META = {
    "name": "categorical_energy",
    "domain": "discrete_energy_sampling",
    "description": "Synthetic categorical-vector classification with local and pairwise energy structure.",
}


def _make(seed, n):
    rng = np.random.default_rng(seed)
    d, k = 8, 3
    x = rng.integers(0, k, size=(n, d))
    unary = np.array([[0.0, 0.45, 0.9], [0.1, 0.0, 0.55], [0.5, 0.15, 0.0]], dtype=np.float32)
    e = unary[x, np.arange(d) % 3].sum(axis=1)
    e += 0.55 * (x[:, :-1] == x[:, 1:]).sum(axis=1)
    e += rng.normal(0.0, 0.18, n)
    y = (e < np.median(e)).astype(np.int64)
    oh = np.eye(k, dtype=np.float32)[x].reshape(n, d * k)
    return oh, y


def get_dataset(seed, n_train, n_test):
    xtr, ytr = _make(int(seed) + 11, n_train)
    xte, yte = _make(int(seed) + 29, n_test)
    return {
        "xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
        "task": "classification", "metric": "err",
        "input_shape": (24,), "out_dim": 2,
    }