import numpy as np

META = {"name": "congestion_moe_sequence", "domain": "moe-routing", "description": "Multi-token regime-conditioned regression where token groups are processed by shared experts."}

def get_dataset(seed, n_train, n_test):
    def make(n, s):
        rng = np.random.default_rng(s)
        length, dim = 12, 4
        regime = rng.integers(0, 4, size=n)
        x = rng.normal(0, 0.7, size=(n, length, dim)).astype(np.float32)
        x[:, :, 0] += (regime[:, None] - 1.5) * 0.8
        x[:, :, 1] += np.sin(regime[:, None] * 1.7)
        y = (0.8 * np.sin(x[:, 0, 0] + regime) + 0.5 * x[:, 3, 1] ** 2
             - 0.35 * x[:, 7, 2] + 0.25 * (regime - 1.5)
             + rng.normal(0, 0.08, n))
        return x, y.astype(np.float32)[:, 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}