import numpy as np

META = {
    "name": "stochastic_intervention_scm",
    "domain": "causal-world-model",
    "description": "Random intervention X causes mediator M and outcome Y; tests joint modeling under intervention shift and mediator selection."
}

def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    # Context is observed but does not explain away the stochastic intervention.
    c_train = rng.randn(n_train, 1).astype("float32")
    c_test = rng.randn(n_test, 1).astype("float32")
    x_train = (rng.randn(n_train, 1) + 0.35*c_train).astype("float32")
    x_test = (2.0*rng.randn(n_test, 1) + 0.35*c_test).astype("float32")
    m_train = (x_train + 0.3*c_train + rng.randn(n_train, 1)).astype("float32")
    m_test = (x_test + 0.3*c_test + rng.randn(n_test, 1)).astype("float32")
    y_train = (x_train + m_train + 0.5*c_train + rng.randn(n_train, 1)).astype("float32")
    y_test = (x_test + m_test + 0.5*c_test + rng.randn(n_test, 1)).astype("float32")
    # x is supplied as an observed feature for training/evaluation; the task target is y.
    return {"xtr": np.concatenate([c_train, x_train, m_train], 1),
            "ytr": y_train[:, 0], "xte": np.concatenate([c_test, x_test, m_test], 1),
            "yte": y_test[:, 0], "task": "regression", "metric": "mse"}