# Some code taken from: https://github.com/wsjeon/maddpg-rllib/

import imp
import os

import gym
from gym import wrappers
from ray import rllib

from multiagent.environment import MultiAgentEnv
import multiagent.scenarios as scenarios


CUSTOM_SCENARIOS = ['simple_switch']


class ParticleEnvRenderWrapper(gym.Wrapper):
    def __init__(self, env, horizon):
        super().__init__(env)
        self.horizon = horizon

    def reset(self):
        self._num_steps = 0

        return self.env.reset()

    def render(self, mode):
        if mode == 'human':
            self.env.render(mode=mode)
        else:
            return self.env.render(mode=mode)[0]

    def step(self, actions):
        obs_list, rew_list, done_list, info_list = self.env.step(actions)

        self._num_steps += 1
        done = (all(done_list) or self._num_steps >= self.horizon)

        # Gym monitor expects reward to be an int.  This is only used for its
        # stats reporter, which we're not interested in.  To make video recording
        # work, we package the rewards in the info object and extract it below.
        return obs_list, 0, done, [rew_list, done_list, info_list]


class RLlibMultiAgentParticleEnv(rllib.MultiAgentEnv):
    def __init__(self, scenario_name, horizon, monitor_enabled=False, video_frequency=500):
        self._env = _make_env(scenario_name, horizon, monitor_enabled, video_frequency)
        self.num_agents = self._env.n
        self.agent_ids = list(range(self.num_agents))

        self.observation_space_dict = self._make_dict(self._env.observation_space)
        self.action_space_dict = self._make_dict(self._env.action_space)

    def reset(self):
        obs_dict = self._make_dict(self._env.reset())
        return obs_dict

    def step(self, action_dict):
        actions = list(action_dict.values())
        obs_list, _, _, infos = self._env.step(actions)
        rew_list, done_list, _ = infos

        obs_dict = self._make_dict(obs_list)
        rew_dict = self._make_dict(rew_list)
        done_dict = self._make_dict(done_list)
        done_dict['__all__'] = all(done_list)
        info_dict = self._make_dict([{'done': done} for done in done_list])

        return obs_dict, rew_dict, done_dict, info_dict

    def render(self, mode='human'):
        self._env.render(mode=mode)

    def _make_dict(self, values):
        return dict(zip(self.agent_ids, values))


def _video_callable(video_frequency):
    def should_record_video(episode_id):
        if episode_id % video_frequency == 0:
            return True
        return False

    return should_record_video


def _make_env(scenario_name, horizon, monitor_enabled, video_frequency):
    if scenario_name in CUSTOM_SCENARIOS:
        # Scenario file must exist locally
        file_path = os.path.join(os.path.dirname(__file__), scenario_name + '.py')
        scenario = imp.load_source('', file_path).Scenario()
    else:
        scenario = scenarios.load(scenario_name + '.py').Scenario()

    world = scenario.make_world()

    env = MultiAgentEnv(world, scenario.reset_world, scenario.reward, scenario.observation)
    env.metadata['video.frames_per_second'] = 8

    env = ParticleEnvRenderWrapper(env, horizon)

    if not monitor_enabled:
        return env

    return wrappers.Monitor(env, './logs/videos', resume=True, video_callable=_video_callable(video_frequency))


def env_creator(config):
    monitor_enabled = False
    if hasattr(config, 'worker_index') and hasattr(config, 'vector_index'):
        monitor_enabled = (config.worker_index == 1 and config.vector_index == 0)

    return RLlibMultiAgentParticleEnv(**config, monitor_enabled=monitor_enabled)