import gymnasium as gym
import numpy as np
import pytest
from gymnasium.envs.mujoco.mujoco_env import MujocoEnv

from open_loop.pd_controller import PDController


@pytest.mark.parametrize(
    "env_tuple",
    [
        ("Swimmer-v4", [-10, 20]),
        ("Swimmer-v4", [30, -20]),
    ],
)
def test_pd_control(env_tuple):
    # Remove time limit
    env_id, desired_qpos = env_tuple
    mujoco_env = gym.make(env_id).unwrapped

    assert isinstance(mujoco_env, MujocoEnv)

    # create a PD controller
    env = PDController(mujoco_env, n_joints=len(desired_qpos), kp=10, kd=0.5)
    desired_qpos = np.deg2rad(desired_qpos)
    for _ in range(100):
        env.step(desired_qpos)

    assert np.allclose(env.current_joint_pos, desired_qpos)
