import pytest

from mlagents_envs.registry import default_registry
from mlagents_envs.side_channel.engine_configuration_channel import (
    EngineConfigurationChannel,
)
from mlagents_envs.base_env import ActionTuple
import numpy as np

BALL_ID = "3DBall"


@pytest.mark.parametrize("n_ports", [1])
def test_set_action_single_agent(base_port: int) -> None:
    engine_config_channel = EngineConfigurationChannel()
    env = default_registry[BALL_ID].make(
        base_port=base_port,
        worker_id=0,
        no_graphics=True,
        side_channels=[engine_config_channel],
    )
    engine_config_channel.set_configuration_parameters(time_scale=100)
    for _ in range(3):
        env.reset()
        behavior_name = list(env.behavior_specs.keys())[0]
        d, t = env.get_steps(behavior_name)
        for _ in range(50):
            for agent_id in d.agent_id:
                action = np.ones((1, 2))
                action_tuple = ActionTuple()
                action_tuple.add_continuous(action)
                env.set_action_for_agent(behavior_name, agent_id, action_tuple)
            env.step()
            d, t = env.get_steps(behavior_name)
    env.close()


@pytest.mark.parametrize("n_ports", [1])
def test_set_action_multi_agent(base_port: int) -> None:
    engine_config_channel = EngineConfigurationChannel()
    env = default_registry[BALL_ID].make(
        base_port=base_port,
        worker_id=0,
        no_graphics=True,
        side_channels=[engine_config_channel],
    )
    engine_config_channel.set_configuration_parameters(time_scale=100)
    for _ in range(3):
        env.reset()
        behavior_name = list(env.behavior_specs.keys())[0]
        d, t = env.get_steps(behavior_name)
        for _ in range(50):
            action = np.ones((len(d), 2))
            action_tuple = ActionTuple()
            action_tuple.add_continuous(action)
            env.set_actions(behavior_name, action_tuple)
            env.step()
            d, t = env.get_steps(behavior_name)
    env.close()