import numpy as np

META = {
    "name": "correlated_token_moe_regression",
    "domain": "moe-routing",
    "description": "Correlated token groups with regime-dependent nonlinear targets for expert routing."
}


def get_dataset(seed, n_train, n_test):
    def make(n, s):
        rng = np.random.RandomState(s)
        x = rng.normal(size=(n, 16, 4)).astype(np.float32)
        regime = (x[:, :, 0].mean(1) > 0).astype(np.float32)
        token = np.where(
            regime[:, None] > 0,
            np.sin(x[:, :, 0]) + 0.45 * x[:, :, 1] ** 2,
            np.cos(x[:, :, 1]) - 0.45 * x[:, :, 0] ** 2,
        )
        y = (
            token.mean(1)
            + 0.25 * x[:, :, 2].mean(1)
            + 0.15 * x[:, :, 3].mean(1)
            + rng.normal(0, 0.06, n)
        ).astype(np.float32)
        return x, y[:, None]

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