import pytest
from ray.rllib.env.multi_agent_env import make_multi_agent
from ray.rllib.tests.test_nested_observation_spaces import NestedMultiAgentEnv


class TestMultiAgentEnv:
    def test_space_in_preferred_format(self):
        env = NestedMultiAgentEnv()
        spaces_in_preferred_format = env._check_if_space_maps_agent_id_to_sub_space()
        assert spaces_in_preferred_format, "Space is not in preferred format"
        env2 = make_multi_agent("CartPole-v1")()
        spaces_in_preferred_format = env2._check_if_space_maps_agent_id_to_sub_space()
        assert (
            not spaces_in_preferred_format
        ), "Space should not be in preferred format but is."

    def test_spaces_sample_contain_in_preferred_format(self):
        env = NestedMultiAgentEnv()
        # this environment has spaces that are in the preferred format
        # for multi-agent environments where the spaces are dict spaces
        # mapping agent-ids to sub-spaces
        obs = env.observation_space_sample()
        assert env.observation_space_contains(
            obs
        ), "Observation space does not contain obs"

        action = env.action_space_sample()
        assert env.action_space_contains(action), "Action space does not contain action"

    def test_spaces_sample_contain_not_in_preferred_format(self):
        env = make_multi_agent("CartPole-v1")({"num_agents": 2})
        # this environment has spaces that are not in the preferred format
        # for multi-agent environments where the spaces not in the preferred
        # format, users must override the observation_space_contains,
        # action_space_contains observation_space_sample,
        # and action_space_sample methods in order to do proper checks
        obs = env.observation_space_sample()
        assert env.observation_space_contains(
            obs
        ), "Observation space does not contain obs"
        action = env.action_space_sample()
        assert env.action_space_contains(action), "Action space does not contain action"


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))