import numpy as np

META = {
    "name": "local_shock_characteristic",
    "domain": "pde",
    "description": "1-D five-variable shock/acoustic local prediction task with equilibrium and oscillatory modes.",
}


def get_dataset(seed, n_train, n_test):
    def make(n, stream):
        rng = np.random.default_rng(int(seed) + stream)
        L = 9
        x = np.linspace(-1.0, 1.0, L)
        xs, ys = [], []
        for _ in range(int(n)):
            center = rng.uniform(-0.35, 0.35)
            width = rng.uniform(0.035, 0.11)
            shock = 0.5 * (1.0 + np.tanh((x - center) / width))
            h = 1.0 + rng.uniform(0.05, 0.35) * shock + 0.012 * rng.normal(size=L)
            q = rng.uniform(-0.18, 0.18) + 0.05 * shock + 0.009 * rng.normal(size=L)
            theta = 1.0 + rng.uniform(-0.15, 0.15) + 0.03 * shock + 0.009 * rng.normal(size=L)
            envelope = np.exp(-((x - center) / rng.uniform(0.18, 0.32)) ** 2)
            amp = rng.uniform(0.03, 0.22)
            ring = ((-1.0) ** np.arange(L)) * envelope
            a4 = amp * ring + 0.003 * rng.normal(size=L)
            a5 = -0.8 * amp * ring + 0.003 * rng.normal(size=L)
            state = np.stack([h, q, theta, a4, a5], axis=1)
            # Stable local one-step target: center value plus a small physical-like flux.
            d = state[1:] - state[:-1]
            target = state[4] - 0.18 * (state[5] - state[3]) + 0.08 * (d[3] - d[2])
            xs.append(state.astype(np.float32).reshape(-1))
            ys.append(np.float32(target[0]))
        return np.asarray(xs, dtype=np.float32), np.asarray(ys, dtype=np.float32)[:, None]

    xtr, ytr = make(n_train, 0)
    xte, yte = make(n_test, 5000)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}