import numpy as np

META = {"name": "latent_mixture_transport", "domain": "diffusion-sampling", "description": "Regression from noisy source Gaussian-mixture states to component-specific transported Gaussian targets."}

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    d = 8
    def make(n):
        z = rng.integers(0, 2, size=n)
        x = rng.normal(size=(n, d)).astype(np.float32)
        x[:, 0] += np.where(z == 0, -2.0, 2.0)
        x[:, 1] += np.where(z == 0, 1.0, -1.0)
        y = np.empty((n, 2), dtype=np.float32)
        y[:, 0] = 0.9*x[:, 0] + 0.25*x[:, 2] + np.where(z == 0, -2.5, 2.5)
        y[:, 1] = -0.65*x[:, 1] + 0.35*x[:, 3] + np.where(z == 0, 1.8, -1.8)
        y += rng.normal(scale=np.where(z[:, None] == 0, 0.65, 1.15), size=(n, 2)).astype(np.float32)
        return x, y
    xtr, ytr = make(n_train)
    xte, yte = make(n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "input_shape": (d,), "out_dim": 2}