import numpy as np

META = {
    "name": "gaussian_score_matching_1d",
    "domain": "loss",
    "description": "Scalar nonlinear Gaussian score-matching task with known N(0,1) target score for Stein discrepancy training."
}

def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    ztr = rng.randn(n_train, 1).astype(np.float32)
    zte = np.random.RandomState(seed + 5000).randn(n_test, 1).astype(np.float32)
    def transport(z):
        return (0.92 * z + 0.12 * np.sin(z)).astype(np.float32)
    return {
        "xtr": ztr, "ytr": transport(ztr),
        "xte": zte, "yte": transport(zte),
        "task": "regression", "metric": "mse", "out_dim": 1
    }