import gym
import envs.envs
from gym_minigrid.wrappers import *
from envs.wrappers import *


OBJ_TYPES = ["key", "ball", "box"]
OBJ_COLORS = ["red", "blue", "green", "purple", "yellow"]

def make_env(env_key="MiniGrid-PickUpObj-v0", goals=None, 
                dist_type=None, dist_color=None, 
                num_dists=4, size=7, seed=None):
    env = gym.make(env_key, size=size, goals=goals, dist_type=dist_type, 
                    dist_color=dist_color, num_dists=num_dists, seed=seed)
    env.seed(seed)
    env = FixEnv(env)
    env = FixRGBImgObsWrapper(env)
    env = ImgObsWrapper(env)
    return env
