import numpy as np

META = {
    'name': 'poisson_action_value',
    'domain': 'dynamics/control',
    'description': 'Continuous two-dimensional action-value regression with smooth isolated quadratic optima.'
}

def get_dataset(seed, n_train, n_test):
    def sample(n, rs):
        state = rs.uniform(-1.0, 1.0, size=(int(n), 2)).astype(np.float32)
        optimum = np.stack([
            0.65 * state[:, 0] + 0.15 * np.sin(2.0 * state[:, 1]),
            -0.55 * state[:, 1] + 0.10 * np.cos(2.0 * state[:, 0]),
        ], axis=1).astype(np.float32)
        action = rs.uniform(-1.0, 1.0, size=(int(n), 2)).astype(np.float32)
        q = -np.sum((action - optimum) ** 2, axis=1, keepdims=True)
        return np.concatenate([state, action], axis=1), q
    xtr, ytr = sample(n_train, np.random.RandomState(int(seed)))
    xte, yte = sample(n_test, np.random.RandomState(int(seed) + 100003))
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'task': 'regression', 'metric': 'mse', 'out_dim': 1}