import numpy as np

META = {
    "name": "manifold_pendulum",
    "domain": "dynamics_and_embedded_manifolds",
    "description": "Controlled pendulum windows with angular state embedded on S1; predict next embedded state and angular velocity.",
}

def get_dataset(seed, n_train, n_test):
    def make(n, s):
        rng = np.random.RandomState(s)
        X = np.empty((n, 24), dtype=np.float32)
        Y = np.empty((n, 3), dtype=np.float32)
        dt = 0.05
        for i in range(n):
            theta = rng.uniform(-np.pi, np.pi)
            omega = rng.uniform(-1.5, 1.5)
            damp = rng.uniform(0.05, 0.25)
            g = rng.uniform(8.5, 11.0)
            rows = []
            for _ in range(9):
                u = rng.uniform(-1.0, 1.0)
                rows.append((theta, omega, u))
                omega = omega + dt * (-g * np.sin(theta) - damp * omega + u)
                theta = theta + dt * omega
            X[i] = np.asarray(rows[:8], dtype=np.float32).reshape(-1)
            th, om, _ = rows[8]
            Y[i] = (np.cos(th), np.sin(th), om)
        return X, Y
    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"}