import numpy as np

META = {
    "name": "budgeted_route_costs",
    "domain": "graph-decision",
    "description": "Two-route parallel graph: predict four edge costs from graph features; decisions are shortest route under budgeted edge delays."
}


def get_dataset(seed, n_train, n_test):
    def make(n, s):
        rng = np.random.RandomState(s)
        z = rng.uniform(0.25, 2.0, (n, 3)).astype(np.float32)
        w, a, b = z.T
        c = np.stack([
            0.75*w + 0.30*a + 0.10*np.sin(2*w),
            0.55*w + 0.25*b + 0.08*np.cos(1.7*w),
            0.70*w + 0.28*b + 0.12*np.sin(1.3*w + a),
            0.72*w + 0.30*a + 0.10*np.cos(1.1*w + b)
        ], 1).astype(np.float32)
        c += rng.normal(0, 0.035, c.shape).astype(np.float32)
        return z, np.maximum(c, 0.03)
    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": 4}