import os
import gym
import h5py
import urllib.request
from d4rl.offline_env import OfflineEnv

DEFAULT_SIZE = 500
DOMAIN = 'antmaze'
TASK = 'multitask-large-play-v0'
TASK_IDX_LIST = list(range(7))

class MultiTaskOfflineEnv(OfflineEnv):
	"""
	An multitask mujoco environment that contains a list of mujoco environments.
	"""
	def __init__(self,
				domain=DOMAIN,
				task=TASK,
				task_idx_list=TASK_IDX_LIST):
		assert (domain is not None and task is not None), (domain, task)

		env_id = f"{domain}-{task}"
		self.mujoco_envs = [gym.envs.make(env_id, task_idx=task_idx) for task_idx in task_idx_list]
		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()

	def get_dataset(self, h5path=None):
		return self.mujoco_envs[0].get_dataset()
