# AUTOGENERATED! DO NOT EDIT! File to edit: 09_CurriculumSchedulerOneMachine.ipynb (unless otherwise specified).

__all__ = ['CurriculumSchedulerOneMachine', 'train_env_specify']

# Cell
import os
import sys
import signal
import subprocess
import time
import numpy as np
import yaml
import random

# Cell
class CurriculumSchedulerOneMachine():
	def __init__(self, yaml_path, command_dic, waiting_second=40):
		"""
		Curriculum Scheduler controls the process of training with multi tasks.
		Args:
			yaml_path (String) : The path of an yaml file which includes all information about task,
								 such as environment, start and goal points pairs, output action scale, threshold ($c$) and window_size ($k$).
			command_dic (Dictionary) : A dictionary maps environment to the `shell command` using to launch the simulation environment.
			waiting_second (Int) : The number of seconds to wait when switching simulation environments.
		"""
		self.command_dic = command_dic
		self.waiting_second = waiting_second
		with open(yaml_path, 'r') as f:
			self.yaml_data = yaml.load(f, Loader=yaml.CLoader)
			self.yaml_sort = sorted(self.yaml_data)
		self.success_queue = []
		self.current_task_index = 0
		self.process = None
		self.__switch_task()

	def update(self, success):
		"""
		Use a buffer and calculate average success rate for performance measure.
		You should update the buffer when the RL trainging episode `reset`.

		Args:
			success (Int) : $1$ if success and $0$ if fail.
		Returns:
			start_goal_pair (Dictionary) : A dictionary about the x,y,z position of start point and goal point.
										   i.e. ""{'goal': [33.65625, -7.6171875, 0.12646484375], 'start': [16.5, 40.40625, 0.12646484375]}""
			action_scale (Dictionary) : A dictionary about the output action scale.
										i.e. ""{'angular': 0.18, 'linear': 0.33, 'ur5': 0.5}"".
		"""
		current_task = self.yaml_sort[self.current_task_index]
		current_task_env = self.yaml_data[self.yaml_sort[self.current_task_index]]["environment"]
		window_size = self.yaml_data[current_task]["window_size"]
		threshold = self.yaml_data[current_task]["threshold"]
		self.success_queue.append(success)
		rate = np.mean(self.success_queue)
		print("------ CurriculumSchedulerOneMachine ------")
		print("Current Task:",current_task,current_task_env)
		print("window_size:",window_size, ", Current Buffer size:",len(self.success_queue))
		print("Threshold:",threshold, ", Current Avg success rate:",rate)

		if len(self.success_queue) >= window_size:
			self.success_queue = self.success_queue[-window_size:]
			if (self.current_task_index<len(self.yaml_sort)-1) and rate>threshold:
				self.current_task_index += 1
				self.success_queue = []
				self.__switch_task()
		start_goal_pair = self.yaml_data[current_task]['pairs'][random.randint(0,len(self.yaml_data[current_task]['pairs'])-1)]
		action_scale = self.yaml_data[current_task]['action_scale']
		return start_goal_pair, action_scale

	def __switch_task(self):
		"""
		Close the old simulation evironment and launch a new one.
		"""
		if not self.process==None:
			os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)  # Send the signal to all the process groups
			for i in range(self.waiting_second):
				time.sleep(1)
				print("waiting", self.waiting_second-i)

		current_task_env = self.yaml_data[self.yaml_sort[self.current_task_index]]["environment"]
		command_str = self.command_dic[current_task_env]
		self.process = subprocess.Popen(command_str, stdout=subprocess.PIPE,shell=True, preexec_fn=os.setsid)
		print('\033[92m@@@@@@@@@ \\\ Level Up /// @@@@@@@@@\033[0m')
		for i in range(self.waiting_second):
			time.sleep(1)
			print("waiting", self.waiting_second-i)


# Cell
def train_env_specify(task_list, data_root='Data'):
	"""
		Automatically generate a well formatted yaml file for the specific environment.

		Args:
			task_list (List) : A List of environment name (String).
							   i.e. ["cave","forest","campus"]
		Returns:
			A yaml file named `curriculum.yaml` would be save at data_root folder.
		"""
	assert os.path.exists(data_root),"data_root Path not found"
	for task in task_list:
		assert os.path.exists(data_root+'/'+task+'.yaml'),"No environment yaml file"

	CL_training = dict()
	for index,task in enumerate(task_list):
		with open(data_root+'/'+task+'.yaml', 'r') as f:
			value = yaml.load(f, Loader=yaml.CLoader)
		key = 'task' + (str(index+1)).zfill(2)
		CL_training[key] = value

	with open(data_root+'/'+'curriculum.yaml', 'w') as f:
		yaml.dump(CL_training, f)