import numpy as np

META = {
    "name": "cayley_cycle_distance",
    "domain": "graph-positional-encoding",
    "description": "Shortest-path distance regression on a cyclic Cayley graph; exact Z_n displacement is the positional coordinate."
}


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    n = 32
    total = int(n_train) + int(n_test)
    u = rng.integers(0, n, size=total)
    v = rng.integers(0, n, size=total)
    delta = (v - u) % n
    y = (np.minimum(delta, n - delta) / (n / 2)).astype(np.float32)
    p = rng.permutation(total)
    tr, te = p[:n_train], p[n_train:]
    return {"xtr": np.stack([u[tr], v[tr]], 1), "ytr": y[tr, None],
            "xte": np.stack([u[te], v[te]], 1), "yte": y[te, None],
            "task": "regression", "metric": "mse", "out_dim": 1, "n_nodes": n}


def encode(ds, kind):
    n = int(ds["n_nodes"])
    def f(x):
        x = np.asarray(x, dtype=np.int64)
        out = np.zeros((len(x), 2 * n), dtype=np.float32)
        rows = np.arange(len(x))
        if kind == "baseline":
            out[rows, x[:, 0]] = 1.0
            out[rows, n + x[:, 1]] = 1.0
        elif kind == "idea":
            # A=Z_n and g_uv=z(v)-z(u) mod n; unused second block keeps dimensions equal.
            d = (x[:, 1] - x[:, 0]) % n
            out[rows, d] = 1.0
        else:
            raise ValueError(kind)
        return out
    out = dict(ds)
    out["xtr"] = f(ds["xtr"])
    out["xte"] = f(ds["xte"])
    out["input_shape"] = (2 * n,)
    return out


def math_check(n=32):
    # Exact quotient/path consistency on the cycle: increments sum to n=0 in Z_n.
    labels = np.arange(n, dtype=np.int64)
    max_edge_error = 0
    for u in range(n):
        v = (u + 1) % n
        max_edge_error = max(max_edge_error, int((labels[v] - labels[u] - 1) % n))
    cycle_sum = int(sum(1 for _ in range(n)) % n)
    return {"group": f"Z_{n}", "cycle_sum_mod_n": cycle_sum,
            "max_edge_increment_error": max_edge_error,
            "path_independence": cycle_sum == 0}