"""
Taken from https://github.com/openai/baselines
"""
import numpy as np
from gym import spaces

from . import VecEnv
from .util import copy_obs_dict, dict_to_obs, obs_space_info


class DummyVecEnv(VecEnv):
    """
    VecEnv that does runs multiple envs sequentially, that is,
    the step and reset commands are send to one environment at a time.
    Useful when debugging and when num_env == 1 (in the latter case,
    avoids communication overhead)
    """

    def __init__(self, env_fns):
        """
        Arguments:

        env_fns: iterable of callables      functions that build envs
        """
        self.envs = [fn() for fn in env_fns]
        env = self.envs[0]
        VecEnv.__init__(self, len(env_fns), env.observation_space, env.action_space)
        self._max_episode_steps = env.unwrapped._max_episode_steps
        obs_space = env.observation_space
        if isinstance(obs_space, spaces.MultiDiscrete):
            obs_space.shape = obs_space.shape[0]

        self.keys, shapes, dtypes = obs_space_info(obs_space)

        self.buf_obs = {k: np.zeros((self.num_envs,) + tuple(shapes[k]), dtype=dtypes[k]) for k in self.keys}
        self.buf_dones = np.zeros((self.num_envs,), dtype=bool)
        self.buf_rews = np.zeros((self.num_envs,), dtype=np.float32)
        self.buf_infos = [{} for _ in range(self.num_envs)]
        self.actions = None

    def step_async(self, actions):
        listify = True
        try:
            if len(actions) == self.num_envs:
                listify = False
        except TypeError:
            pass

        if not listify:
            self.actions = actions
        else:
            assert self.num_envs == 1, "actions {} is either not a list or has a wrong size - cannot match to {} envs".format(
                actions, self.num_envs)
            self.actions = [actions]

    def step_wait(self):
        for e in range(self.num_envs):
            action = self.actions[e]
            if isinstance(self.envs[e].action_space, spaces.Discrete):
                action = int(action)

            obs, self.buf_rews[e], self.buf_dones[e], self.buf_infos[e] = self.envs[e].step(action)
            self._save_obs(e, obs)
        return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones),
                self.buf_infos.copy())

    def reset_mdp(self):
        for e in range(self.num_envs):
            obs = self.envs[e].reset_mdp()
            self._save_obs(e, obs)
        return self._obs_from_buf()

    def reset(self, task=None):
        for e in range(self.num_envs):
            if task is not None:
                obs = self.envs[e].reset(task=task[e])
            else:
                obs = self.envs[e].reset()
            self._save_obs(e, obs)
        return self._obs_from_buf()

    def _save_obs(self, e, obs):
        for k in self.keys:
            if k is None:
                self.buf_obs[k][e] = obs
            else:
                self.buf_obs[k][e] = obs[k]

    def _obs_from_buf(self):
        return dict_to_obs(copy_obs_dict(self.buf_obs))

    def get_images(self):
        return [env.render(mode='rgb_array') for env in self.envs]

    def render(self, mode='human'):
        if self.num_envs == 1:
            self.envs[0].render(mode=mode)
        else:
            super().render(mode=mode)

    def get_task(self):
        return self.envs[0].unwrapped.get_task()

    def get_belief(self):
        return self.envs[0].unwrapped.get_belief()

    def reset_task(self, task=None):
        if self.num_envs == 1:
            return self.envs[0].unwrapped.reset_task(task)
        else:
            raise NotImplementedError

    def get_env_attr(self, attr):
        return getattr(self.envs[0].unwrapped, attr)
