import numpy as np

META = {
    "name": "phase_synchronization_regression",
    "domain": "geometry/graph synchronization",
    "description": "Multi-view rotated node vectors with noisy pairwise relative-phase observations; predict a gauge-invariant latent regression target.",
}
N = 8
EDGES = [(i, j) for i in range(N) for j in range(i + 1, N)]


def _make(seed, n):
    rng = np.random.default_rng(seed)
    theta = rng.uniform(-np.pi, np.pi, N)
    theta -= theta[0]
    phi = np.array([theta[i] - theta[j] + rng.normal(0.0, 0.28) for i, j in EDGES], dtype=np.float32)
    x = np.empty((n, N, 2), dtype=np.float32)
    y = np.empty(n, dtype=np.float32)
    c, s = np.cos(theta).astype(np.float32), np.sin(theta).astype(np.float32)
    for k in range(n):
        u = rng.normal(0.0, 1.0, (N, 2)).astype(np.float32)
        x[k, :, 0] = c * u[:, 0] - s * u[:, 1]
        x[k, :, 1] = s * u[:, 0] + c * u[:, 1]
        y[k] = np.tanh(u[:, 0].sum() / 3.0) + 0.25 * np.sin(u[:, 1].sum())
    obs = np.tile(phi, (n, 1)).astype(np.float32)
    return np.concatenate([x.reshape(n, -1), obs], axis=1), y[:, None]


def get_dataset(seed, n_train, n_test):
    xtr, ytr = _make(seed, n_train)
    xte, yte = _make(seed + 5000, n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}