import numpy as np

META = {
    "name": "obb_clearance_trajectory",
    "domain": "geometric_collision_avoidance",
    "description": "Eight-step ego trajectory regression with an observed oriented obstacle box."
}

def get_dataset(seed, n_train, n_test):
    def make(n, s):
        r = np.random.RandomState(s)
        start_y = r.uniform(-.12, .12, n)
        obs_x = r.uniform(1.8, 2.2, n)
        obs_y = r.uniform(-.18, .18, n)
        obs_yaw = r.uniform(-.25, .25, n)
        x = np.stack([start_y, obs_x, obs_y, obs_yaw], 1).astype(np.float32)
        t = np.linspace(.5, 4., 8, dtype=np.float32)[None, :]
        yy = start_y[:, None] + (obs_y-start_y)[:, None] * t/4
        y = np.stack([np.broadcast_to(t, (n,8)), yy], -1)
        y += r.normal(0, .015, y.shape).astype(np.float32)
        return x, y.reshape(n,16).astype(np.float32)
    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":16}