import os
import inspect
from collections import defaultdict, OrderedDict
import types
import dill
import threading
from enum import Enum
import logging

from .rmscore import RMSEntryType, RMSEntry, Pipe, Task, Resource, UnrunTask, VirtualModule
from. rmscore import RMSUpdateEvent
from .rmscore import _get_func_ba, _has_resource_content
from .mphelper import ProcessWrap, ProcessWrapState

_logger = logging.getLogger("rmspool")


class RMSUnrunTasksBuilder():
	'''	
	Create a virtual RMS system that groups the pipeline into unruntasks. 
	These unruntasks can then be used on the RMS pool
	'''
	def __init__(self, rmspool):
		self.rmspool = rmspool
		self.rms = rmspool.rms
		self.rms.registerRMSUpdateListener(self)
		
		self.vfdb = {}
		self.unruntasks = []
		self.vm = VirtualModule()
		
	def onRMSUpdate(self, events):
		'''
		Builder allows users to specify a virtual file (that is generated by some tasks)
		
		The builder receives update from RMS and try to replace these virtual futures.
		
		# To be done: unregister when the vfdb is empty. 
		'''
		for event in events:
			if event[0] == RMSUpdateEvent.INSERT:
				if event[1][0] is RMSEntryType.FILERESOURCE:
					f = self.rms.get(event[1])
					if f.file_path in self.vfdb:
						vr = self.vfdb.pop(f.file_path)
						self.rms.replace_virtualresource(vr, f)
	@property
	def sql(self):
		return self.rms.sql
	
	def register_file(self, *args, **kwargs):
		return self.rms.register_file(*args, **kwargs)
	
	def file_from_path(self, file_path):
		try:
			return self.rms.file_from_path(file_path)
		except:
			file_path = os.path.abspath(file_path)
			if file_path in self.vfdb:
				return self.vfdb[file_path]
			else:
				vr = self.rms.create_virtualresource()
				self.vfdb[file_path] = vr
				return vr
	
	def register_pipe(self, *args, **kwargs):
		pipe = self.rms.register_pipe(*args, **kwargs)
		# Create an alternative pipe with self stored as rmsp 
		pipe_rb = Pipe(pipe.pid, pipe.func, pipe.return_volatile, pipe.is_deterministic, pipe.module_name, pipe.func_name, pipe.output_func, pipe.description, pipe.tags, pipe.info, self)
		return pipe_rb
	
	def get_pipe(self, *args, **kwargs):
		pipe = self.rms.get_pipe(*args, **kwargs)
		pipe_rb = Pipe(pipe.pid, pipe.func, pipe.return_volatile, pipe.is_deterministic, pipe.module_name, pipe.func_name, pipe.output_func, pipe.description, pipe.tags, pipe.info, self)
		return pipe_rb
	
	def rmsimport(self, namespace=None, moduleobj=None, varname = None, altname=None, return_first_module=False):
		'''
		Perform a virtual import
	
		'''
		def is_bound_method(func):
			return inspect.ismethod(func)
		pipe_results = self.sql.cursor().execute(f"SELECT pid FROM pipes;").fetchall()
		pipes = [self.get_pipe(pid) for pid, in pipe_results]
		
		if namespace is None:
			namespace = dict()
		if moduleobj is None:
			moduleobj = "builtins"

		# There are two ways to add bound_method. 
		# 1. If the class is already registered, set the bound method as the class attr
		# 2. If the class is not registered, set the class as another virtual module 
		bound_method_classes = defaultdict(list)
		for p in pipes:
			if inspect.ismethod(p.func):
				if inspect.isclass(p.func.__self__):
					bound_method_classes[p.func.__self__].append(p)
		
		
		# Import the module
		if varname is None:
			output_varnames = []
			output_pipes = []
			for pipe in pipes:
				if pipe.func.__module__ == moduleobj and not inspect.ismethod(pipe.func): #(inspect.isclass(pipe.func) or inspect.isfunction(pipe.func)):
					output_varnames.append(pipe.func.__name__)
					output_pipes.append(pipe)
			d = {v:p for v, p in zip(output_varnames, output_pipes)}
			if altname is None:
				vm = self.vm
				for i in moduleobj.split("."):
					if not hasattr(vm, i):
						setattr(vm, i, VirtualModule())
					vm = getattr(vm, i)
				for v, p in d.items():
					setattr(vm, v, p)
				
				# Deal with bound method that is not in output pipes
				for k, bound_methods in bound_method_classes.items():
					if k.__module__ in moduleobj:
						for p in output_pipes:
							if k is p.func:
								break
						else:
							i = k.__name__
							if not hasattr(vm, i):
								setattr(vm, i, VirtualModule())
							for bound_method in bound_methods:
								setattr(getattr(vm, i), bound_method.func.__name__, bound_method)
				first_vm_name = moduleobj.split(".")[0]
				namespace[first_vm_name] = getattr(self.vm, first_vm_name)
				if return_first_module:
					returnvalue = namespace[first_vm_name]
				else:
					returnvalue = vm
			else:
				# Deal with bound method that is not in output pipes
				to_add = {}
				for k, bound_methods in bound_method_classes.items():
					if k.__module__ in moduleobj:
						for p in output_pipes:
							if k is p.func:
								break
						else:
							i = k.__name__
							if i not in to_add:
								to_add[i] = VirtualModule()
							for bound_method in bound_methods:
								setattr(to_add[i], bound_method.func.__name__, bound_method)
				namespace[altname] = types.SimpleNamespace(**d, **to_add)
				returnvalue = namespace[altname]

				
		# Import the variable
		else:
			output_pipes = []
			# wild card case
			if varname == "*":
				tmp_returnvalues = {}
				for pipe in pipes:
					if pipe.func.__module__ == moduleobj and not inspect.ismethod(pipe.func):#(inspect.isclass(pipe.func) or inspect.isfunction(pipe.func)):
						namespace[pipe.func.__name__] = pipe
						output_pipes.append(pipe)
						tmp_returnvalues[pipe.func.__name__] = namespace[pipe.func.__name__]
				to_add = {}
				for k, bound_methods in bound_method_classes.items():
					if k.__module__ in moduleobj:
						for p in output_pipes:
							if k is p.func:
								break
						else:
							i = k.__name__
							if i not in to_add:
								to_add[i] = VirtualModule()
							for bound_method in bound_methods:
								setattr(to_add[i], bound_method.func.__name__, bound_method)
				for k, p in to_add.items():
					namespace[k] = p
					tmp_returnvalues[k] = namespace[k]
				if len(tmp_returnvalues) < 1:
					raise Exception(f"Cannot find anything to import in {moduleobj}")
				returnvalue = tmp_returnvalues
				
			# normal case
			else:
				tmp_returnvalues = [] 
				for pipe in pipes:
					if pipe.func.__module__ == moduleobj and not inspect.ismethod(pipe.func):#(inspect.isclass(pipe.func) or inspect.isfunction(pipe.func)):
						if pipe.func.__name__ == varname:
							if altname is None:
								altname = varname
							namespace[altname] = pipe
							output_pipes.append(pipe)
							tmp_returnvalues.append(pipe)
				to_add = {}
				
				for k, bound_methods in bound_method_classes.items():
					if k.__name__ == varname and k.__module__ in moduleobj:
						for p in output_pipes:
							if k is p.func:
								break
						else:
							i = k.__name__
							if i not in to_add:
								to_add[i] = VirtualModule()
							for bound_method in bound_methods:
								setattr(to_add[i], bound_method.func.__name__, bound_method)
				for k, p in to_add.items():
					namespace[k] = p
					tmp_returnvalues.append(namespace[k])
				if len(tmp_returnvalues) < 1:
					raise Exception(f"Cannot find {varname} in {moduleobj}")
				if len(tmp_returnvalues) > 1:
					assert False
				returnvalue = tmp_returnvalues[0]
				
		for pipe in output_pipes:
			if inspect.isclass(pipe.func):
				if pipe.func in bound_method_classes:
					for bound_method in bound_method_classes[pipe.func]:
						setattr(pipe, bound_method.func.__name__, bound_method)
							
		return returnvalue
	
	def run(self, pipe, args=(), kwargs={},
				task_description="", task_tags=set(), task_info={},
				resource_description="", resource_tags=set(), resource_info={}, 
				fileresource_description="", fileresource_tags=set(), fileresource_info={}):
		'''
		A virtual run method. Rather than running the RMS pipe, it creates unruntasks and virtualresources. 
		'''
		vr = self.rms.create_virtualresource()
		if self.rms.scriptID is not None:
			task_info = {"scriptid":self.rms.scriptID, **task_info}
		
		u = self.rms.create_unruntask(pipe.pid, args, kwargs, [vr], None,
									task_description, task_tags, task_info,
									resource_description, resource_tags, resource_info,
									fileresource_description, fileresource_tags, fileresource_info
									)
		self.unruntasks.append(u)
		return vr

	def execute_builder(self):
		self.rmspool.run(*self.unruntasks)
		l = list(self.unruntasks)
		self.unruntasks.clear()
		return l


class RMSPoolUpdateEvent(Enum):
	SUBMISSION = 1
	EXECUTION = 2
	COMPLETION = 3
	SKIPPED = 4
	ERROR = 5
	CANCELLATION = 6
	


def _dill_load(path):
	with open(path, "rb") as fr:
		return dill.load(fr)
class RMSProcessWrapPool():
	'''
	Process is not recycled. Hence this is very inefficient for running many short processes. 
	However, it benefits from interruptable Process, and not affected by jupyter's "stop"   
	'''
	def __init__(self, rms, nthread, use_thread=False): #, use_hybrid_process_thread=False
		self.rms = rms
		
		self.nthread = nthread
		
		self.pending_tasks_lock = threading.Condition() # A lock for accessing pending_tasks and next_pid
		self.pending_tasks = OrderedDict()
		self.next_pid = 0
		
		self.finished_tasks_lock = threading.Condition() # A lock for accessing finished_tasks and error_tasks
		self.finished_tasks = set()
		self.error_tasks = set()
		
		self.futures_lock = threading.Condition() # A lock for accessing futures 
		self.futures = OrderedDict()

		self.resources_to_fetch = set()
		self.fetch_resource_tasks = OrderedDict()
		
		self.update_event_lock = threading.Condition()
		self.update_events = list()

		self.closing = False
		
		self.use_thread = use_thread
		# The thread run tills closing is True and no more pending tasks remain.
		threading.Thread(target=self._run_pending_tasks).start()
		threading.Thread(target=self._fire_update_events).start()
		
		
		self.listeners = []
		
		self.rms.registerRMSUpdateListener(self)
		
	def registerListener(self, listener):
		self.listeners.append(listener)
		
	def fireUpdateEvent(self, events):
		self.update_event_lock.acquire()
		self.update_events.extend(events)
		self.update_event_lock.notify_all()
		self.update_event_lock.release()
		
		
	def _fire_update_events(self):
		
		self.update_event_lock.acquire()
		while not self.closing:
			while len(self.update_events) > 0:
				events = list(self.update_events)
				self.update_events.clear()
				self.update_event_lock.release() # Release the lock here so that we won't be bothered too much on handling to on update event
				
				for listener in self.listeners:
					try:
						listener.onProcessPoolUpdate(events)
					except Exception as e:
						print(e)
					
				self.update_event_lock.acquire() # Reacquire the lock
			self.update_event_lock.wait()
			
		self.update_event_lock.release()
	
	def onRMSUpdate(self, events):
		'''
		An unruntask is pending if it still has virtual resources not yet resolved.
		Hence, upon any update of the virtual resource from RMS, 
		try to notify pending_tasks_lock to check if any unruntasks can now be submitted   
		'''
		toNotify = False
		for event in events:
			if event[0] == RMSUpdateEvent.DELETE:
				if event[1][0] == RMSEntryType.VIRTUALRESOURCE:
					toNotify = True
		if toNotify:
			self.pending_tasks_lock.acquire()
			self.pending_tasks_lock.notify_all()
			self.pending_tasks_lock.release()
			
			
	def run(self, *unruntasks, priority=20):
		'''
		Submit one or more unruntask to the pool. An unrun task will start running if all of its dependencies are resolved. 
		'''
		if self.closing:
			raise ValueError("Invalid state")
		self.pending_tasks_lock.acquire()
		self._replace_unruntasks_virtualresource(unruntasks)
		registered_pids = []
		for unruntask in unruntasks:
			pid = unruntask.uid # This pid is process id, not pipe id
			self.pending_tasks[pid] = unruntask, priority
			registered_pids.append(pid)
		self.fireUpdateEvent([(RMSPoolUpdateEvent.SUBMISSION, unruntask.uid) for unruntask in unruntasks])
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		
		for unruntask in unruntasks:
			self.fetch_resource_content(*unruntask.input_resources, priority=priority)
		return registered_pids
	
	def fetch_resource_content(self, *resources, priority=20):
		if self.closing:
			raise ValueError("Invalid state")
		self.pending_tasks_lock.acquire()
		self._fetch_resource_content(*resources, priority=priority)
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		
	def get_total_running_processes(self):
		self.futures_lock.acquire()
		self.finished_tasks_lock.acquire()
		total_running_processes = len(self.futures) - len(self.finished_tasks) - len(self.error_tasks)
		self.finished_tasks_lock.release()
		self.futures_lock.release()
		return total_running_processes

				
				
	def _fetch_resource_content(self, *resources, priority=20, force_fetching=False, force_fetching_ids=None):
		# Before calling this, acquire the pending_tasks lock
		resources = list(resources)
		while (len(resources) > 0):
			resource = resources.pop()
			if not resource.has_content:
				# Check if there exists a pre-saved resource
				if (not force_fetching or (force_fetching_ids is not None and resource.rid not in force_fetching_ids)) and _has_resource_content(self.rms.resource_dump_dir, resource):
					if resource.rid in self.futures:
						continue
					if resource.rid not in self.pending_tasks or self.pending_tasks[resource.rid][1] > priority:
						self.pending_tasks[resource.rid] = (resource, priority)
				else:
					if resource.tid is None:
						raise Exception("Fail. Unable to retrieve resource tid.")
					task = self.rms.get_task(resource.tid)
					pipe = self.rms.get_pipe(task.pid)
					
					if not pipe.is_deterministic:
						raise Exception("Fail. Unable to rerun pipe because it is a non-deterministic pipe.")
					if not len(task.output_files) == 0:
						raise Exception("Fail. Re-retrieving the content may overwrite certain files.")
					if any(not os.path.exists(f.file_path) for f in task.input_fileresources):
						raise Exception("Fail. Some files are missing")
					# Add the task with priority
					# Maybe I need to add future lock here
					if task.tid in self.futures:
						continue
					if not task.tid in self.pending_tasks or self.pending_tasks[task.tid][1] > priority:
						self.pending_tasks[task.tid] = (task, priority)
						resources.extend(task.input_resources)

	def _replace_unruntasks_virtualresource(self, unruntasks):
		'''
		When users submit an unruntask that uses virtual resources with replacement, replace them. 
		'''
		for unruntask in unruntasks:
			for virtualresource in unruntask.input_virtualresources:
				if virtualresource.replacement is not None:
					resource = virtualresource.replacement
					original_kv = list(unruntask.ba.arguments.items())
					for k, v in original_kv:
						if unruntask.ba.signature.parameters[k].kind == inspect.Parameter.VAR_POSITIONAL:
							unruntask.ba.arguments[k] = [resource if (isinstance(i, RMSEntry) and i is virtualresource) else i for i in v]
						elif unruntask.ba.signature.parameters[k].kind == inspect.Parameter.VAR_KEYWORD:
							unruntask.ba.arguments[k] = {ik:(resource if (isinstance(i, RMSEntry) and i is virtualresource) else i) for ik, i in v.items()}
						else:
							if isinstance(v, RMSEntry):
								if v is virtualresource:
									unruntask.ba.arguments[k] = resource
				
	def _run_pending_tasks(self):
		'''
		Check pending tasks and submit them to the pool executor if the dependencies are fulfilled.
		
		(1) Check for skippable unruntasks
		(2) Update required resources to be fetched
		
		The submitted pending tasks are removed from self.pending_tasks
		
		This thread will be run at the beginning, until all pending tasks are complete and close
		'''
		
		self.pending_tasks_lock.acquire()
		while not self.closing or len(self.pending_tasks) > 0: # Loop until closing and no more pending tasks remain
			# Check if I can run any jobs from the pending tasks
			total_running_processes = self.get_total_running_processes()
			if total_running_processes < self.nthread:
				nidlethread = self.nthread - total_running_processes
			
				# Check for tasks that have been done before. If so, skip it.
				_logger.info(f"Checking tasks to skip...")
				skipped = []
				for pid, (job, priority) in self.pending_tasks.items():
					if isinstance(job, UnrunTask):
						unruntask = job
						if len(unruntask.input_virtualresources) == 0: # If an unruntask still have virtual resource, we cannot tell whether it has been run before
							pipe = self.rms.get_pipe(unruntask.pid)
							ba = unruntask.ba
							prev_tasks = self.rms.find_tasks_by_pipe_and_args(pipe.pid, ba.args, ba.kwargs)
							if len(prev_tasks) > 0:
								skipped.append(unruntask.uid)
								task = prev_tasks[0]
								self.rms.replace_unruntask(unruntask, task)
				if len(skipped) > 0:
					for pid in skipped:
						self.pending_tasks.pop(pid)
					print(f"Skipped {len(skipped)} tasks that were done before.")
					self.fireUpdateEvent([(RMSPoolUpdateEvent.SKIPPED, uid) for uid in skipped] + [(RMSPoolUpdateEvent.EXECUTION, uid) for uid in skipped])
					continue # Repeat from the beginning, because there may be other skippable tasks
				
				# If no unruntask is done before, check if they need certain resource. If so, load the fetching process into the queue 
				# Update pending tasks by looking at required resource fetch
				tmp = list(self.pending_tasks.items()) # Avoid issue when updating pending_tasks				
				for pid, (job, priority) in tmp:
					if isinstance(job, UnrunTask):
						unruntask = job
						if len(unruntask.input_virtualresources) == 0:
							if not all(r.has_content for r in unruntask.input_resources):
								self._fetch_resource_content(*unruntask.input_resources, priority = priority)

				# Now check for all the pending tasks available to run
				_logger.info(f"Checking pending tasks to submit...")
				available_to_run_pids = []
				for pid, (job, priority) in self.pending_tasks.items():
					if isinstance(job, UnrunTask):
						unruntask = job
						if len(unruntask.input_virtualresources) == 0:
							if all(r.has_content for r in unruntask.input_resources):
								available_to_run_pids.append(pid)
					elif isinstance(job, Task):
						if all(r.has_content for r in job.input_resources):
							available_to_run_pids.append(pid)
					elif isinstance(job, Resource):
						# It should be runnable any time
						available_to_run_pids.append(job.rid)
					else:
						raise Exception("Job must be Task or Unruntask or Resource")
				_logger.info(f"Total pending pids available to submit: {len(available_to_run_pids)} / {len(self.pending_tasks)}")
				
				if len(available_to_run_pids) > 0:
					
					available_to_run_pids.sort(key=lambda pid: self.pending_tasks[pid][1]) # Sort by priority
					executed = []
					available_to_run_pid_idx = 0
					while nidlethread > 0 and available_to_run_pid_idx < len(available_to_run_pids):
						pid = available_to_run_pids[available_to_run_pid_idx]
						available_to_run_pid_idx += 1
						job, priority = self.pending_tasks.pop(pid)
						try:
							if isinstance(job, UnrunTask):
								# Perform a new task
								unruntask = job
								pipe = self.rms.get_pipe(unruntask.pid)
								ba = unruntask.ba
								
								executed.append(unruntask.uid)
								func_args, func_kwargs = self.rms._ba_rms_to_func_args_kwargs(ba)
								pwcallback=(lambda state, raw_return_value, begin_time, end_time, pid, callback=None,unruntask=unruntask:
																			self._completion_callback(state, raw_return_value, begin_time, end_time, pid, callback, unruntask))
								
								# Attempts to use process first
								use_thread = self.use_thread
								if pipe.return_volatile:
									use_thread = True
								pw = ProcessWrap(pipe.func, func_args, func_kwargs, pwcallback, pid, use_thread=use_thread)
								pw.start()
													
								self.futures_lock.acquire()
								self.futures[pid] = pw
								self.futures_lock.release()
								nidlethread -= 1
							elif isinstance(job, Task):
								# Fetch resource from existing tasks
								task = job
								pipe = self.rms.get_pipe(task.pid)
								tid = task.tid
								
								executed.append(tid)
								ba = _get_func_ba(pipe.func, task.args, task.kwargs)
								func_args, func_kwargs = self.rms._ba_rms_to_func_args_kwargs(ba)
								
								use_thread = self.use_thread
								if pipe.return_volatile:
									use_thread = True
								pwcallback=(lambda state, raw_return_value, begin_time, end_time, tid, callback=None,task=task:
									self._resource_fetch_completion_callback(state, raw_return_value, begin_time, end_time, tid, callback, task))
								
								pw = ProcessWrap(pipe.func, func_args, func_kwargs, pwcallback, tid, use_thread=use_thread)
								pw.start()
								
								self.futures_lock.acquire()
								self.futures[pid] = pw
								self.futures_lock.release()
								nidlethread -= 1
								
							elif isinstance(job, Resource):
								resource = job
								rid = resource.rid
								use_thread = self.use_thread
								pwcallback=(lambda state, raw_return_value, begin_time, end_time, rid, resource=resource, priority=priority:
									self._resource_load_completion_callback(state, raw_return_value, rid, resource, priority))
								
								pw = ProcessWrap(_dill_load, [self.rms.resource_dump_dir + "/" + rid], {}, pwcallback, rid, use_thread=use_thread)
								pw.start()
								
								self.futures_lock.acquire()
								self.futures[pid] = pw
								self.futures_lock.release()
								nidlethread -= 1
						except Exception as e:
							print(f"WARNING! An unruntask cannot be started with exceptions! pid: {pid}")
							print(str(e))
						
					_logger.info(f"Remaining pending tasks: {len(self.pending_tasks)}")
				else:
					self.pending_tasks_lock.wait()
			else:
				self.pending_tasks_lock.wait()
		self.pending_tasks_lock.release()
	
	def _resource_load_completion_callback(self, state, raw_return_value, pid, resource, priority):
		self.pending_tasks_lock.acquire() # The lock is important since the callback may occassionally triggers run_pending_tasks
		self.futures_lock.acquire() # I don't think I need this here
		self.finished_tasks_lock.acquire()
		try:
			if self.futures[pid].state == ProcessWrapState.COMPLETE:
				self.fireUpdateEvent([(RMSPoolUpdateEvent.COMPLETION, pid)])
				self.finished_tasks.add(pid)
				resource.content = raw_return_value
			elif self.futures[pid].state == ProcessWrapState.ERROR:
				# If error occurs (e.g. the dumped file is broken), try to fetch resource content instead of loading
				print(f"WARNING: Cannot load resource {pid} from file. Try to fetch the resource instead.")
				self._fetch_resource_content(resource, priority=priority, force_fetching=True, force_fetching_ids=[resource.rid])
				self.error_tasks.add(pid)
		except Exception as e:
			print(f"WARNING: Something happens when loading the resource! rid: {pid}")
			print(str(e))
			
		self.finished_tasks_lock.notify_all()
		self.finished_tasks_lock.release()
		self.futures_lock.release()
		
		
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		
	def _resource_fetch_completion_callback(self, state, raw_return_value, begin_time, end_time, pid, callback, task):
		self.pending_tasks_lock.acquire() # The lock is important since the callback may occassionally triggers run_pending_tasks
		self.futures_lock.acquire() # I don't think I need this here, but maybe needed if need to deal with the pickling error
		self.finished_tasks_lock.acquire()
		tid = pid
		try:
			if self.futures[pid].state == ProcessWrapState.COMPLETE:
				self.fireUpdateEvent([(RMSPoolUpdateEvent.COMPLETION, pid)])
				self.finished_tasks.add(pid)
				resource = task.output_resources[0]
				resource.content = raw_return_value
			elif self.futures[pid].state == ProcessWrapState.PICKLING_ERROR:
				pipe = self.rms.get_pipe(task.pid)
				ba = _get_func_ba(pipe.func, task.args, task.kwargs)
				func_args, func_kwargs = self.rms._ba_rms_to_func_args_kwargs(ba)
				
				pwcallback=(lambda state, raw_return_value, begin_time, end_time, tid, callback=None,task=task:
					self._resource_fetch_completion_callback(state, raw_return_value, begin_time, end_time, tid, callback, task))
				
				pw = ProcessWrap(pipe.func, func_args, func_kwargs, pwcallback, tid, use_thread=True)
				pw.start()
				
				self.futures[pid] = pw
				
				
			else:
				self.fireUpdateEvent([(RMSPoolUpdateEvent.ERROR, pid)])
				self.error_tasks.add(pid)
				# Handle errors: Stop all running tasks?
				pass
		except Exception as e:
			print(f"WARNING: Something happens when registering the finished tasks! pid: {pid}")
			print(str(e))
			
		self.finished_tasks_lock.notify_all()
		self.finished_tasks_lock.release()
		self.futures_lock.release()
		
		
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()

		
	def _completion_callback(self, state, raw_return_value, begin_time, end_time, pid, callback, unruntask):
		'''
		Add pids to finished tasks. 
		Notify pending tasks
		'''
		self.pending_tasks_lock.acquire() # The lock is important since the callback may occassionally triggers run_pending_tasks
		self.futures_lock.acquire() # I don't think I need this here
		self.finished_tasks_lock.acquire()
		try:
			if self.futures[pid].state == ProcessWrapState.COMPLETE:
				self.fireUpdateEvent([(RMSPoolUpdateEvent.COMPLETION, pid)])
				self.finished_tasks.add(pid)
				
				pipe = self.rms.get_pipe(unruntask.pid)
				task = self.rms.register_finished_task(raw_return_value, begin_time, end_time, unruntask.ba, pipe, 
													unruntask.task_description, unruntask.task_tags, unruntask.task_info,
													unruntask.resource_description, unruntask.resource_tags, unruntask.resource_info, 
													unruntask.fileresource_description, unruntask.fileresource_tags, unruntask.fileresource_info)
								
				self.rms.replace_unruntask(unruntask, task)
			elif self.futures[pid].state == ProcessWrapState.PICKLING_ERROR:
				# Try to rerun the process by using thread
				pipe = self.rms.get_pipe(unruntask.pid)
				ba = unruntask.ba
				func_args, func_kwargs = self.rms._ba_rms_to_func_args_kwargs(ba)
				pwcallback=(lambda state, raw_return_value, begin_time, end_time, pid, callback=None,unruntask=unruntask:
				self._completion_callback(state, raw_return_value, begin_time, end_time, pid, callback, unruntask))
				
				pw = ProcessWrap(pipe.func, func_args, func_kwargs, pwcallback, pid, use_thread=True)
				pw.start()
				self.futures[pid] = pw
			else:
				self.fireUpdateEvent([(RMSPoolUpdateEvent.ERROR, pid)])
				self.error_tasks.add(pid)
				# Handle errors: Stop all running tasks?
				pass
		except Exception as e:
			import traceback
			traceback.print_exc()
			print(f"WARNING: Something happens when registering the finished tasks! pid: {pid}")
			print(str(e))
		self.finished_tasks_lock.notify_all()
		self.finished_tasks_lock.release()
		self.futures_lock.release()
		
		
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
	


	def cancel(self, pids):
		'''
		Attempts to cancel a task if it is still pending. 
		If a task is being executed or completed, it will be cancelled.  
		See also kill. 
		'''
		removed_from_pending_tasks = []
		self.pending_tasks_lock.acquire()
		for pid in pids:
			if pid in self.pending_tasks:
				self.pending_tasks.pop(pid)
				removed_from_pending_tasks.append(pid)
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		
# 		logging.info(f"Removed from pending tasks: {join_str(removed_from_pending_tasks, ',')}")
# 		logging.info(f"Removed from pools: {join_str(removed_from_futures, ',')}")
	def cancel_all(self):
		'''
		Attempts to cancel all pending tasks 
		'''
		removed_from_pending_tasks = []
		self.pending_tasks_lock.acquire()
		for pid in list(self.pending_tasks.keys()):
			self.pending_tasks.pop(pid)
			removed_from_pending_tasks.append(pid)
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		
		
	def kill(self, pids):
		self.pending_tasks_lock.acquire()
		self.futures_lock.acquire()
		for pid in pids:
			self.futures[pid].kill()
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		self.futures_lock.release()
		
	def kill_all(self):
		'''
		Attempts to kill all running tasks
		'''
		self.pending_tasks_lock.acquire()
		self.futures_lock.acquire()
		for pid in list(self.futures.keys()):
			self.futures[pid].kill()
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		self.futures_lock.release()

# 		logging.info(f"Removed from pending tasks: {join_str(removed_from_pending_tasks, ',')}")
# 		logging.info(f"Removed from pools: {join_str(removed_from_futures, ',')}")
	def clear(self):
		'''
		'''
		raise NotImplementedError()
		pass
	
	def close(self):
		'''
		Terminate the MultiProcessRun
		'''
		self.closing = True
		
		# Notify the pending_tasks_lock so that the thread knows to terminate
		self.pending_tasks_lock.acquire()
		self.pending_tasks_lock.notify_all()
		self.pending_tasks_lock.release()
		
		self.get()
		#self.pool.shutdown()
		

	def __enter__(self): 
		return self
	
	def __exit__(self, type, value, traceback): 
		self.close()
	
	
