from typing import cast
import os
import pytest
from xml.etree.ElementTree import fromstring

from omegaconf import OmegaConf

import pymasep as pm
from test.unit.mocks.game_mock import GameMock
from test.unit.mocks.controller_mock import ControllerMock
from test.unit.mocks.none_action_controller_mock import NoneActionControllerMock


@pytest.fixture(scope="function")
def relative_path():
    yield '../../'


@pytest.fixture(scope='function')
def configuration(base_directory_unit):
    cfg = OmegaConf.load(os.path.join(base_directory_unit, 'config', 'game_one_agent', 'app.yaml'))
    cfg['root_path'] = base_directory_unit
    yield cfg


def test_init():
    env = pm.common.Environment()

    assert isinstance(env, pm.common.Environment)
    assert isinstance(env.agents, set)
    assert env.agents == set()
    assert env.controllers == dict()
    assert not env.end_episode
    assert env.current_step == 0
    assert env.game is None
    assert env.current_state is None
    assert env.current_additional_info == dict()
    assert env.next_action == {}


def test_get_observation_for_agent():
    env = pm.common.Environment()
    env.game = GameMock()
    env.current_state = env.game.init_state(env)
    expected_observation = env.current_state

    real_observation = env.get_observation_for_agent()

    assert expected_observation == real_observation


def test_calculate_end_episode_not_end(base_directory_unit, mock_environment_conf):
    mock_environment_conf.current_state = mock_environment_conf.game.init_state(mock_environment_conf)

    mock_environment_conf.calculate_end_episode()
    is_episode_terminated = mock_environment_conf.end_episode

    assert not is_episode_terminated


def test_calculate_next_state(mock_environment):
    mock_environment.current_state = mock_environment.game.initialize_state(mock_environment)
    mock_environment.choose_action()

    mock_environment.calculate_next_state()

    assert mock_environment.current_step == 1
    assert isinstance(mock_environment.current_state, pm.common.State)
    assert len(mock_environment.current_state.objects) == 3
    assert len(mock_environment.current_state.agents) == 1
    assert mock_environment.current_state.objects['State.object1'].object_state.characteristics['Step0'].value == 1
    assert mock_environment.current_additional_info == {'info': 'Step=Step+1'}


def test_calculate_reward(mock_environment):
    mock_environment.current_state = mock_environment.game.initialize_state(mock_environment)
    mock_environment.choose_action()
    mock_environment.calculate_next_state()

    mock_environment.calculate_reward()

    assert mock_environment.current_additional_info['reward'] == {'State.agent1': 1}


def test_create_base_object_base_object_created(mock_environment):
    obj = mock_environment.create_base_object(name='object1',
                                              template=mock_environment.game.templates['DefaultObjectTemplate'])

    assert isinstance(obj, pm.common.Object)
    assert obj.parent is None


def test_create_base_object_base_object_created_with_parent(mock_environment):
    obj_parent = mock_environment.create_base_object(name='object1',
                                                     template=mock_environment.game.templates['DefaultObjectTemplate'])
    obj = mock_environment.create_base_object(name='object1',
                                              template=mock_environment.game.templates['DefaultObjectTemplate'],
                                              parent=obj_parent)
    assert isinstance(obj, pm.common.Object)
    assert obj_parent == obj.parent


def test_create_agent(mock_environment):
    expected_name = 'AgentName1'
    obj_parent = mock_environment.create_base_object('State',
                                                     template=mock_environment.game.templates['SystemStateTemplate'])

    ag = mock_environment.create_object(name=expected_name,
                                        template=mock_environment.game.templates['DefaultAgentTemplate'],
                                        xml_node=None,
                                        parent=obj_parent)
    ag = cast(pm.common.Agent, ag)

    assert len(mock_environment.agents) == 1
    assert list(mock_environment.agents)[0] == ag
    assert ag.get_fullname() == 'State.AgentName1'
    assert ag.name == expected_name
    assert isinstance(ag.controller, pm.common.Controller)
    assert mock_environment.game.get_system_value(obj_parent, 'AgentOrder').data == [ag.get_fullname()]


def test_create_object(mock_environment):
    expected_name = 'ObjectName1'
    object_tmpl = mock_environment.game.templates['EmptyObjectTemplate']

    obj = mock_environment.create_object(name=expected_name,
                                         template=object_tmpl)

    assert len(mock_environment.agents) == 0
    assert isinstance(obj, pm.common.Object)
    assert obj.name == expected_name


def test_create_object_object_created_from_xml(mock_environment):
    object_xml_str = '<Object name="object1" template="DefaultObjectTemplate" initializer="objectinitializer">' \
                     '<ObjectState name="state">' \
                     '<Characteristic name="Step0"><int>1</int></Characteristic>' \
                     '</ObjectState>' \
                     '</Object>'
    xml_node = fromstring(object_xml_str)

    obj = mock_environment.create_object(name='',
                                         xml_node=xml_node)

    assert len(mock_environment.agents) == 0
    assert len(obj.object_state.characteristics) == 3
    assert isinstance(obj, pm.common.Object)
    assert obj.name == 'object1'


def test_create_action(mock_environment):
    ac = mock_environment.create_action(GameMock.ACTION_TYPE_2)

    assert ac.environment == mock_environment
    assert ac.type == GameMock.ACTION_TYPE_2


def test_create_action_from_xml(mock_environment):
    xml_node_str = '<Action type="1" />'
    xml_node = fromstring(xml_node_str)

    ac = mock_environment.create_action(xml_node=xml_node)

    assert ac.environment == mock_environment
    assert ac.type == GameMock.ACTION_TYPE_1


def test_choose_action_waiting_all(mock_environment):
    mock_environment.game.coord_method = mock_environment.game.MULTIPLAYER_COORDINATION_WAITING_ALL
    mock_environment.current_state = mock_environment.game.initialize_state(mock_environment)
    ag1 = cast(pm.common.Agent, mock_environment.current_state.objects['State.agent1'])
    ag2 = mock_environment.create_object('agent2',
                                         template=mock_environment.game.templates['DefaultAgentTemplate'],
                                         parent=mock_environment.current_state)

    mock_environment.choose_action()

    # agents have chosen its action
    assert ag1.action.type == GameMock.ACTION_TYPE_2
    assert ag2.action.type == GameMock.ACTION_TYPE_2
    # the join action is inside the next environment action
    assert ag1.action == mock_environment.next_action['State.agent1']
    assert ag2.action == mock_environment.next_action['State.agent2']


def test_choose_action_waiting_turn(mock_environment):
    mock_environment.current_state = mock_environment.game.initialize_state(mock_environment)
    ag1 = cast(pm.common.Agent, mock_environment.current_state.objects['State.agent1'])
    ag2 = mock_environment.create_object('agent2',
                                         template=mock_environment.game.templates['DefaultAgentTemplate'],
                                         parent=mock_environment.current_state)

    mock_environment.choose_action()

    # agent1 has chosen its action
    assert ag1.action.type == GameMock.ACTION_TYPE_2
    # the action of ag1 is inside the next environment action
    assert mock_environment.next_action['State.agent1'] == ag1.action
    # no other agent choose its action
    assert ag2.name not in mock_environment.next_action


def test_choose_action_waiting_turn_action_none(mock_environment):
    mock_environment.current_state = mock_environment.game.initialize_state(mock_environment)
    ag1 = cast(pm.common.Agent, mock_environment.current_state.objects['State.agent1'])
    ag1.controller = NoneActionControllerMock(environment=mock_environment)

    mock_environment.choose_action()

    # agent1 has chosen its action but its None
    assert ag1.action is None
    # no action in the next environment action
    assert len(mock_environment.next_action) == 0


def test_clear_all_state_elements_all_clear_except_controllers(mock_environment):
    _ = mock_environment.create_object('agent1',
                                       template=mock_environment.game.templates['DefaultAgentTemplate'])
    mock_environment.current_step = 10
    mock_environment.end_episode = True

    mock_environment.clear_all_state_elements()

    assert len(mock_environment.agents) == 0
    assert len(mock_environment.next_action) == 0
    assert mock_environment.current_state is None
    assert mock_environment.current_step == 0
    assert len(mock_environment.current_additional_info) == 0
    assert not mock_environment.end_episode
    assert len(mock_environment.controllers) == 1
    assert mock_environment.game is not None


def test_get_controller_new_controller(mock_environment):
    agent_name = 'agent1'

    result = mock_environment.get_controller(agent_name, ControllerMock)

    assert isinstance(result, ControllerMock)
    assert len(mock_environment.controllers) == 1
    assert mock_environment.controllers[agent_name] == result


def test_get_controller_existing_controller(mock_environment):
    agent_name = 'agent1'

    ag1 = mock_environment.create_object(agent_name,
                                         template=mock_environment.game.templates['DefaultAgentTemplate'])

    result = mock_environment.get_controller(agent_name, ControllerMock)

    assert isinstance(result, ControllerMock)
    assert len(mock_environment.controllers) == 1
    assert mock_environment.controllers[agent_name] == result
    assert ag1.controller == result
