"""Custom manifold-valued one-step dynamics benchmark."""
import numpy as np

META = {
    "name": "sphere_one_step_dynamics",
    "domain": "geometric_dynamics",
    "description": "Predict one geodesic step of a controlled smooth vector field on S2."
}


def _exp(x, v):
    r = np.linalg.norm(v, axis=1, keepdims=True)
    return np.cos(r) * x + np.sinc(r / np.pi) * v


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)

    def make(n):
        z = rng.uniform(-0.85, 0.85, n)
        angle = rng.uniform(0, 2 * np.pi, n)
        q = np.sqrt(1 - z * z)
        x = np.column_stack([q * np.cos(angle), q * np.sin(angle), z]).astype(np.float32)
        c = np.array([0.35, -0.42, 0.28], dtype=np.float32)
        b = c[None, :] - x * (x * c[None, :]).sum(1, keepdims=True)
        y = _exp(x, 0.12 * b).astype(np.float32)
        return x, y

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