import numpy as np

META = {
    "name": "minkowski_composition",
    "domain": "structured_convex_composition",
    "description": "Regression from two direction-indexed convex latent tuples to their Minkowski-additive composition."
}

N_DIR = 8
D = 2
ANGLES = np.linspace(0.0, 2.0 * np.pi, N_DIR, endpoint=False)


def _tuple(rng):
    center = rng.uniform(-1.0, 1.0, size=2)
    radii = rng.uniform(0.25, 1.25, size=2)
    phase = rng.uniform(-0.35, 0.35)
    ang = ANGLES + phase
    return center[None, :] + np.stack([radii[0] * np.cos(ang), radii[1] * np.sin(ang)], axis=1)


def _make(seed, n):
    rng = np.random.RandomState(seed)
    x = np.empty((n, 2 * N_DIR * D), dtype=np.float32)
    y = np.empty((n, N_DIR * D), dtype=np.float32)
    for i in range(n):
        a, b = _tuple(rng), _tuple(rng)
        x[i] = np.concatenate([a.reshape(-1), b.reshape(-1)])
        y[i] = (a + b).reshape(-1)
    return x, y


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": N_DIR * D}