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


DEFAULT_SIZE = 500
# DOMAIN = ['HalfCheetahVel', 'HalfCheetahVelJump', 'HalfCheetahVelBackward', 'HalfCheetahVelBackwardJump']
# DOMAIN = ['HalfCheetahVel', 'HalfCheetahJump', 'HalfCheetahVelBackward']
DOMAIN = ['Walker2d', 'Walker2dBackward', 'Walker2dJump']
# DOMAIN = ['Walker2d', 'Walker2dJump']
TASK = 'v2'

class MultiTaskEnv(gym.Env):
	"""
	An multitask mujoco environment that contains a list of mujoco environments.
	"""
	def __init__(self,
				domains=DOMAIN,
				task=TASK):
		assert (domains is not None and task is not None), (domains, task)

		self.mujoco_envs = []
		for i, domain in enumerate(domains):
			task_one_hot_id = np.zeros(len(DOMAIN))
			task_one_hot_id[i] = 1.0
			env_id = f"{domain}-{task}"
			self.mujoco_envs.append(gym.envs.make(env_id, task_one_hot_id=task_one_hot_id))
			print('Adding task {} with task id {}'.format(DOMAIN[i], self.mujoco_envs[-1].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, width=width, height=height, depth=depth)

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