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

__all__ = ['CurriculumSchedulerMultiMachine']

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

# Cell
class CurriculumSchedulerMultiMachine():
	def __init__(self, yaml_path):
		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.current_task_env = ''
		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("------ CurriculumSchedulerMultiMachine ------")
		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):
		"""
		A flag for Gym to decide which socket package should be received.
		"""
		self.current_task_env = self.yaml_data[self.yaml_sort[self.current_task_index]]["environment"]
