import numpy as np

META = {
    "name": "sparse_cycle_compatibility",
    "domain": "masked_categorical_compatibility",
    "description": "Paired conditional categorical prediction on a sparse bipartite support with a compatible-joint generator."
}
N = 8
EDGES = np.array([(0,0),(0,1),(1,1),(1,2),(2,2),(2,3),(3,3),(3,4),
                  (4,4),(4,5),(5,5),(5,6),(6,6),(6,7),(7,7),(7,0),
                  (0,4),(2,6),(4,0),(6,2)], dtype=np.int64)

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    ux, vy = rng.normal(0, .45, N), rng.normal(0, .45, N)
    joint_logits = np.array([ux[x] + vy[y] for x, y in EDGES])
    p = np.exp(joint_logits - joint_logits.max()); p /= p.sum()
    def sample(k):
        z = rng.choice(len(EDGES), k, p=p)
        return EDGES[z, 1].astype(np.int64), EDGES[z, 0].astype(np.int64)
    ytr, xtr_target = sample(n_train)
    yte, xte_target = sample(n_test)
    return {
        "xtr": np.eye(N, dtype=np.float32)[ytr], "ytr": xtr_target,
        "xte": np.eye(N, dtype=np.float32)[yte], "yte": xte_target,
        "task": "classification", "metric": "err", "out_dim": N,
        "n": N, "edges": EDGES.copy(), "train_y": ytr, "test_y": yte,
        "train_x": xtr_target, "test_x": xte_target,
    }