import os

from gym import error, spaces
from gym.utils import seeding
import numpy as np
from os import path
from collections import deque
import gym
from gym.spaces import Box

from metaworld.envs.mujoco.sawyer_xyz.v2.sawyer_multitask_drawer_open_v2 import SawyerMultitaskDrawerOpenEnvV2
from metaworld.envs.mujoco.sawyer_xyz.v2.sawyer_multitask_drawer_close_v2 import SawyerMultitaskDrawerCloseEnvV2
from metaworld.envs.mujoco.sawyer_xyz.v2.sawyer_multitask_door_v2 import SawyerMultitaskDoorEnvV2
from metaworld.envs.mujoco.sawyer_xyz.v2.sawyer_multitask_door_close_v2 import SawyerMultitaskDoorCloseEnvV2

DEFAULT_SIZE = 500
TASKS = [SawyerMultitaskDoorEnvV2, SawyerMultitaskDoorCloseEnvV2, SawyerMultitaskDrawerOpenEnvV2, SawyerMultitaskDrawerCloseEnvV2]

class MultiTaskMujocoEnv(gym.Env):
	"""
	An multitask mujoco environment that contains a list of mujoco environments.
	"""
	def __init__(self,
				tasks=TASKS):
		self.mujoco_envs = []
		for i, env in enumerate(TASKS):
			task_one_hot_id = np.zeros(len(TASKS))
			task_one_hot_id[i] = 1.0
			self.mujoco_envs.append(env(task_one_hot_id=task_one_hot_id))
		self.task_idx = 0
		self.observation_space = self.mujoco_envs[self.task_idx].observation_space
		self.action_space = self.mujoco_envs[self.task_idx].action_space

	def seed(self, seed=None):
		self.np_random, seed = seeding.np_random(seed)
		return [seed]

	def step(self, action):
		obs, reward, done, info = self.mujoco_envs[self.task_idx].step(action)
		return obs, reward, done, info

	def reset(self):
		self.task_idx = (self.task_idx + 1) % len(self.mujoco_envs)
		obs = self.mujoco_envs[self.task_idx].reset()
		return obs

	@property
	def dt(self):
		return self.mujoco_envs[self.task_idx].model.opt.timestep * self.mujoco_envs[self.task_idx].frame_skip

	def render(self, mode='human', width=DEFAULT_SIZE, height=DEFAULT_SIZE, depth=False):
		return self.mujoco_envs[self.task_idx].render(mode=mode)

	def close(self):
		self.mujoco_envs[self.task_idx].close()
