"""Custom structurally matched unordered point-set denoising track."""
import numpy as np

META = {
    "name": "unordered_pointset_denoising",
    "domain": "point-set-diffusion",
    "description": "Denoising noisy unordered one-dimensional point configurations; sets are sorted only for canonical representation."
}


def _split(seed, n):
    rng = np.random.RandomState(seed)
    # Repulsive configurations: sorted Gaussian locations, with a mild global
    # scale variation. The evaluation target is the clean set itself.
    z = np.sort(rng.normal(size=(n, 6)).astype(np.float32), axis=1)
    z = z / (np.std(z, axis=1, keepdims=True) + 0.35)
    noisy = z + rng.normal(0, 0.22, size=z.shape).astype(np.float32)
    noisy.sort(axis=1)
    return noisy.astype(np.float32), z.astype(np.float32)


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