import numpy as np

META = {
    "name": "router_regime_regression",
    "domain": "moe-routing",
    "description": "Regression with three observable nonlinear regimes, requiring a learned fixed-budget expert router."
}


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    def make(n):
        x = rng.uniform(-2.0, 2.0, size=(n, 4)).astype(np.float32)
        # Three regimes are observable from x but have distinct functions.
        z = np.argmax(np.stack([x[:, 0], -x[:, 0], x[:, 1]], axis=1), axis=1)
        y = np.where(z == 0, np.sin(2*x[:, 1]) + 0.25*x[:, 2],
            np.where(z == 1, 0.5*x[:, 0]**2 - x[:, 2],
                     x[:, 1]*x[:, 2] + 0.4*x[:, 3]))
        y += 0.04 * rng.normal(size=n)
        return x, y.astype(np.float32)[:, None]
    xtr, ytr = make(n_train); xte, yte = make(n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}