"""
Helpers for constructing vector environments.
"""
from environments.atari.vec_env_utils import Monitor
from environments.empty_room.gym_maze import MazeEnv

from vec_env import DummyVecEnv, ShmemVecEnv


def make_env(rank, seed, env_kwargs):
    env = MazeEnv(
        **env_kwargs
    )
    env = Monitor(env, allow_early_resets=True)
    return env


def make_vec_env(num_env, seed, env_kwargs={}):
    """
    Create a wrapped, monitored SubprocVecEnv for Atari and MuJoCo.
    """
    def make_thunk(rank):
        return lambda: make_env(
            rank=rank,
            seed=seed,
            env_kwargs=env_kwargs,
        )

    if num_env > 1:
        return ShmemVecEnv([make_thunk(i) for i in range(num_env)])
    else:
        return DummyVecEnv([make_thunk(i) for i in range(num_env)])


if __name__ == '__main__':
    env = make_vec_env(4, 0, env_kwargs={'size': 7})
    print(env.num_envs)
    print(env.observation_space)
    print(env.action_space)
    ob = env.reset()
    print(ob.shape)
