"""
# @Time    : 2021/7/1 8:44 上午
# @Author  : hezhiqiang01
# @Email   : hezhiqiang01@baidu.com
# @File    : env_wrappers.py
Modified from OpenAI Baselines code to work with multi-agent envs
"""

import numpy as np


def _stack_dict_obs(obs_list):
    """将 list[dict[str, ndarray]] 堆叠为 dict[str, ndarray]。"""
    keys = obs_list[0].keys()
    return {k: np.array([o[k] for o in obs_list]) for k in keys}


# single env
class DummyVecEnv():
    def __init__(self, env_fns):
        self.envs = [fn() for fn in env_fns]
        env = self.envs[0]
        self.num_envs = len(env_fns)
        self.observation_space = env.observation_space
        self.share_observation_space = env.share_observation_space
        self.action_space = env.action_space
        self.actions = []

    def step(self, actions, episode=-1):
        """
        Step the environments synchronously.
        This is available for backwards compatibility.
        """
        self.step_async(actions)
        return self.step_wait(episode)

    def step_async(self, actions):
        self.actions = actions

    def step_wait(self, episode):
        # ① 与所有子环境交互
        results = [env.step(a) for (a, env) in zip(self.actions, self.envs)]
        # results = [(obs, rew, done, info), ...]
        obs, rews, dones, infos, active_mask, bad_mask, available_actions = map(list, zip(*results))   # 用 list 保持元素可变

        # ② 逐线程检查 done→reset
        for i, done in enumerate(dones):
            done_flag = done if isinstance(done, (bool, np.bool_)) else np.all(done)
            if done_flag:
                obs_i, active_i, available_i = self.envs[i].reset(episode)
                obs[i]   = obs_i
                active_mask[i] = active_i
                available_actions[i] = available_i

        # ③ 转回便于上层使用的格式
        obs  = _stack_dict_obs(obs)           # dict[str, ndarray]
        rews = np.array(rews)
        dones = np.array(dones)
        active_mask = np.array(active_mask)
        bad_mask = np.array(bad_mask)
        available_actions = np.array(available_actions)
        # infos 保持 list[dict] —— Runner 需要这样
        return obs, rews, dones, infos, active_mask, bad_mask, available_actions


    def reset(self, episode):
        result = [env.reset(episode) for env in self.envs]
        obs_list, active_masks, available_actions = zip(*result)        # unzip
        obs = _stack_dict_obs(obs_list)                                 # dict[str, ndarray]
        return obs, np.stack(active_masks), np.stack(available_actions)

    def get_global_state(self):
        return _stack_dict_obs([env.get_global_state() for env in self.envs])

    def close(self):
        for env in self.envs:
            env.close()

    def render(self, mode="human"):
        if mode == "rgb_array":
            return np.array([env.render(mode=mode) for env in self.envs])
        elif mode == "human":
            for env in self.envs:
                env.render(mode=mode)
        else:
            raise NotImplementedError