import sys
import re
import copy

from py_autoflow.program import Program

class Batch:
	general_computation_attrib = { 
		'default': {
			'cpu' : None, 
			'mem' : None, 
			'time' : None,
			'node' : None,
			'multinode' : None,
			'ntask' : None,
			'additional_job_options' : None
		}
	}
	@classmethod
	def init_indexes(cls): # To clean start job relations indexes
		cls.all_batch = {}
		cls.batch_iterator_relations = {}
		cls.jobs_names = []
		cls.nested_iteration_relations = {}
		Program.all_jobs_relations = {} # Reset job index

	def __init__(self, tag, init, main_command, id, exec_folder):
		self.set_initialize(tag, init, main_command, id, exec_folder)

	def set_initialize(self, tag = '', init = '', main_command = '', id = '', exec_folder = None): # method to initialize and to clean the object
		self.regex_deps = None
		tag, init, main_command = self.replace_regexp(tag, init, main_command)
		self.name = None
		self.id = id
		self.iterator = [None] 
		self.parent = None
		self.dependencies = [] # [batch_name, dependency_type, keyword2replace, dep_info] (dep_info is a regex math only used with nested dependencies)
			# None => There isn't dependecies, 
			# 'simple' => One job needs a previous job, 
			# '1to1' => A job in a batch needs another job in other batch, 
			# '*to1' => A job need a previous full batch of jobs
			# 'local' => A simple job needs one job of a batch
		self.initialization = init
		self.main_command = main_command
		self.attrib = {**Batch.general_computation_attrib['default'], **{
			'done' : False,
			'folder' : True,
			'buffer' : False,
			'exec_folder' : exec_folder,
			'cpu_asign' : None} # number, list or mono
		}
		self.get_name_and_iterators_and_modifiers(tag)
		self.set_execution_attrib()
		self.set_cpu()
		self.jobs = []
		Batch.all_batch[self.name] = self

	def clean(self):
		self.set_initialize()

	def replace_regexp(self, tag, init, main_command):
		if 'JobRegExp:' in tag: tag = self.scan_JobRegExp_tag(tag)
		init = self.resolve_regexp(init)
		main_command = self.resolve_regexp(main_command)
		return tag, init, main_command

	def resolve_regexp(self, instructions):
		if isinstance(instructions, str):
			while '!JobRegExp:' in instructions: instructions = self.scan_JobRegExp(instructions)
		return instructions

	def scan_JobRegExp(self, command):
		data = re.search(r'!JobRegExp:([^ \n]+):([^ \n]+)!([^ \n]+)', command) # *to1 with regexp
		#data.group(0) => reference string (command), 
		#data.group(1) => batch_pattern, 
		#data.group(2) => iterator_pattern, 
		#data.group(3) => adyacent string to regexp as regexp/file_name
		job_names = self.get_dependencies_by_regexp(data.group(1), data.group(2))
		if len(job_names) > 0: self.regex_deps = 'command'
		new_string = ' '.join([jn + ')' + data.group(3) for jn in job_names ])
		command = command.replace(data.group(0), new_string)
		return command

	def scan_JobRegExp_tag(self, tag):
		data = re.search(r'JobRegExp:([^ \n]+):([^;\] \n]+)', tag) # 1to1 with regexp
		#data.group(0) => reference string (command), 
		#data.group(1) => batch_pattern, 
		#data.group(2) => iterator_pattern
		job_names = self.get_dependencies_by_regexp(data.group(1), data.group(2))
		if len(job_names) > 0: self.regex_deps = 'tag'
		new_string = ';'.join([jn + ')' for jn in job_names ])
		tag = tag.replace(data[0], new_string)
		return tag

	def get_dependencies_by_regexp(self, batch_pattern, iterator_pattern):
		selected_batches = [ cmd_name for cmd_name in Batch.all_batch.keys() if re.search(batch_pattern, cmd_name) ]
		job_names = []
		for batch_name in selected_batches:
			iterators = Batch.all_batch[batch_name].iterator
			if iterators[0] != None: iterators = [ re.sub(r'[&!%)]','',it) for it in iterators ]
			if iterator_pattern != '-':
				if iterators[0] == None: continue 
				iterators = [ iter for iter in iterators if re.search(batch_pattern, iter) ]
			if len(iterators) > 0:
				if iterators[0] == None: 
					job_names.append(batch_name)
				else:
					for iter in iterators: job_names.append(batch_name+iter)		
		return job_names

	def get_name_and_iterators_and_modifiers(self, tag):
		match_tag = re.search(r"(^.+)\[([^\]]+)\]\)", tag) # iterative node
		if match_tag == None: # Non iterative node (simple node)
			match_tag = re.search(r"(^.+)\)", tag)
		name = match_tag.groups()[0]
		self.name , self.attrib['done'], self.attrib['folder'], self.attrib['buffer'] = self.check_execution_modifiers(name)
		if len(match_tag.groups()) == 2:
			self.iterator = []
			for interval in match_tag.groups()[1].split(';'):
				if '-' in interval:
					limits = interval.split('-')
					self.iterator = self.iterator + [ str(i) for i in range(int(limits[0]),int(limits[1])) ]
				else:
					self.iterator.append(interval)
		Batch.batch_iterator_relations[self.name] = self.iterator

	def check_execution_modifiers(self, name, iter_type = False): #The last paremeter is used to indicate tha name is a iterator not an orignal node name
		done = False
		folder = True
		buffer = False
		if '%' in name: done = True
		if '!' in name: folder = False
		if '&' in name: buffer = True
		if not iter_type:
			name = re.sub(r"&|!|%|\)", '', name) # Delete function characters
		else:
			name = re.sub(r"&|!|%", '', name) # Delete function characters
		return name, done, folder, buffer

	def set_execution_attrib(self):
		if self.initialization != None: self.initialization = self.scan_resources(self.initialization)
		if isinstance(self.main_command, str): self.main_command = self.scan_resources(self.main_command) 

	def scan_resources(self, command):
		resources_line = None
		for line in command.splitlines():
			if 'resources:' in line:
				line = line.rstrip()
				resources_line = line
				fields = line.split(' ')
				for index, field in enumerate(fields):
					if field == '-r':
						self.attrib.update(Batch.general_computation_attrib[fields[index+1]])
					elif field == '-c':
						self.attrib['cpu'] = int(fields[index+1])
					elif field == '-m':
						self.attrib['mem'] = fields[index+1]
					elif field == '-n':
						self.attrib['node'] = fields[index+1]
					elif field == '-t':
						self.attrib['time'] = fields[index+1]
					elif field == '-u':
						self.attrib['multinode'] = int(fields[index+1])
					elif field == '-A':
						self.attrib['additional_job_options'] = [ pair.split('=') for pair in fields[index+1].split(';') ]

				if '-s' in fields:
					self.attrib['ntask'] = True
				else:
					self.attrib['ntask'] = False
		if resources_line != None: command = command.replace(resources_line, '')
		return command

	def set_cpu(self):
		if self.initialization != None: self.initialization = self.scan_cpu(self.initialization) 
		if isinstance(self.main_command, str): self.main_command = self.scan_cpu(self.main_command)
		if self.attrib['cpu_asign'] == 'mono': self.attrib['cpu'] = 1 

	def scan_cpu(self, command):
		if '[cpu]' in command:
			command = command.replace('[cpu]', str(self.attrib['cpu']))
			self.attrib['cpu_asign'] = 'number'
		elif '[lcpu]' in command:
			command = command.replace('[lcpu]', 'workers')
			self.attrib['cpu_asign'] = 'list'
		elif self.attrib.get('cpu_asign') == None:
			self.attrib['cpu_asign'] = 'mono'
		return command

	# def handle_dependencies(self, dinamic_variables):
	# 	for instructions in [self.initialization, self.main_command]:
	# 		if isinstance(instructions, str):
	# 			#self.scan_dependencies(instructions) # NOT NECESSARY? REMOVED BY COLLISION CON REGEX SYSTEM. THE DINAMYC VARIABLES ARE NO USED
	# 			dinamic_variables = dinamic_variables + self.collect_dinamic_variables(instructions)
	# 			self.dependencies = self.dependencies + self.check_dependencies_with_DinVar(instructions, dinamic_variables)
	# 	return dinamic_variables

	def asign_child_batch(self):
		batches = []
		if isinstance(self.main_command, list):
			for id in self.main_command:
				batch = self.get_batch(id)
				batch.parent = self.name
				batches.append(batch)
			self.main_command = batches

	def get_batch(self, id):
		selected_batch = None
		for name, batch in Batch.all_batch.items():
			if batch.id == id:
				selected_batch = batch  
				break
		return selected_batch

	def has_jobs(self):
		res = len(self.jobs) > 0
		return res
	

	def get_jobs(self):
		jobs = []
		if isinstance(self.main_command, list): # There are nested batchs
			#raise Exception("Detected nested tasks, there are not implemented")
			self.add_nested_jobs(jobs)
		else:
			self.check_regex_dependencies() # Not implemented implement with nested tasks
			for num, it in enumerate(self.iterator):
				job_attrib = copy.deepcopy(self.attrib)
				if not it == None:
					it, done, job_attrib['folder'], job_attrib['buffer'] = self.check_execution_modifiers(it, True)
					if not self.attrib['done']: job_attrib['done'] = done  # To keep attrib priority in batch on job
				if it == None:
					it_str = ''
				else:
					it_str = it
				name = f"{self.name}{it_str}"
				job_dependencies = []
				batch_deps = len(self.dependencies)
				initialization = self.replace_dependencies(self.initialization, job_dependencies, it, num)
				parameters = self.replace_dependencies(self.main_command, job_dependencies, it, num)
				self.dependencies = self.dependencies[0:len(self.dependencies) - (batch_deps+1)] # Clean temporal dependencies by regexp
				job_dependencies = list(set(job_dependencies))
				job = Program(name, initialization, parameters, job_dependencies, job_attrib)
				job.batch = self.name
				Batch.jobs_names.append(job.name)
				jobs.append(job)
				self.jobs.append(job)
		return jobs

	def replace_dependencies(self, command, job_dependencies, it, num):
		if not command == None:
			if isinstance(command, str): command = command.replace('(*)', str(it)) 
			self.scan_dependencies(command)
			new_string = None
			for batch_name, dep_type, dep_keyword2replace, dep_info in self.dependencies:
				if dep_type == 'simple':
					if Batch.all_batch[batch_name].parent == None: 	new_string = batch_name + ')'
					job_dependencies.append(batch_name)
				elif dep_type == '1to1':
					if Batch.all_batch[batch_name].parent == None or not self.parent == None:
						dep_name = f"{batch_name}{Batch.batch_iterator_relations[batch_name][num]}"
					else:
						root_batch = self.get_root(batch_name)
						selected_jobs = root_batch.get_jobs_by_batch_name(batch_name)
						dep_name = selected_jobs[num].name
					job_dependencies.append(dep_name)
					new_string = dep_name + ')' 
				elif dep_type == '*to1':
					if Batch.all_batch[batch_name].parent == None or not self.parent == None:
						def get_dep(it):
							dep_name =  batch_name + it
							job_dependencies.append(dep_name)
							return f"{dep_name}){dep_info}"
						new_string = ' '.join(list(map(get_dep, Batch.batch_iterator_relations[batch_name])))
					else:
						root_batch = self.get_root(batch_name)
						selected_jobs = root_batch.get_jobs_by_batch_name(batch_name)
						def get_dep(j):
							job_dependencies.append(j.name)
							return f"{j.name + ')'}{dep_info}"
						new_string = ' '.join(list(map(get_dep, selected_jobs)))
					dep_keyword2replace = f"{dep_keyword2replace}{dep_info}"
				elif dep_type == 'local':
					if Batch.all_batch[batch_name].parent == None or not self.parent == None:
						if dep_info in list(map(lambda it: it.replace(')',''), Batch.batch_iterator_relations[batch_name])): #This avoids cross dependencies by similar names, map used for regexp deps
							dep_name = batch_name + dep_info
							job_dependencies.append(dep_name)
							new_string = dep_name + ')'
					else:
						dep_name = dep_keyword2replace.replace(')','') #This avoids cross dependencies by similar names
						job_dependencies.append(dep_name)
						new_string = dep_name + ')'

				if not dep_keyword2replace == None and not new_string == None: command = command.replace(dep_keyword2replace, new_string) 
		return command

	def scan_dependencies(self, command):
		if not command == None: # When command is the initialize, sometimes can be undefined
			batches = []
			for k, val in Batch.all_batch.items(): #sorting is used to match last jobs first, and avoid small matches of first nodes
				batches.append([k , val])
			matched_regions = []
			for name, batch in reversed(batches):

				if name+')' in command and not self.string_overlap(matched_regions, name+')', command):
					self.dependencies.append([name, 'simple', name+')', None])
				if f"!{name}*!" in command and not self.string_overlap(matched_regions, f"!{name}*!", command):
					self.dependencies.append([name, '1to1', f"!{name}*!", None])
				if f"!{name}!" in command and not self.string_overlap(matched_regions, f"!{name}!", command):
					for string_match in re.findall('!' + name + r'!([^ \n]+)', command):
						self.dependencies.append([name, '*to1', f"!{name}!", string_match])
				local_dependencies = re.findall(name + r'([^\( \n]+)\)', command)
				for local_dependency in local_dependencies:
					if not self.string_overlap(matched_regions, f"{name}{local_dependency}"+')', command):
						self.dependencies.append([name, 'local', f"{name}{local_dependency}"+')', local_dependency])

	def string_overlap(self, matched_regions, substr, string):
		match = False
		str_range = self.get_string_position(substr, string)
		if len(str_range) > 0:
			for start, ending in matched_regions:
				if ((str_range[0] >= start and str_range[0] <= ending) or 
					(str_range[-1] >= start and str_range[-1] <= ending) or
					(str_range[0] <= start and str_range[-1] >= ending)):
					match = True
					break
			matched_regions.append(str_range)
		return match

	def get_string_position(self, substr, string):
		start = string.index(substr)
		ending = start + len(substr) - 1
		str_range = [start, ending]
		return str_range

	def add_nested_jobs(self, jobs):
		temp_jobs = []
		for batch in self.main_command:
			js = batch.get_jobs()
			temp_jobs.extend(js)
		jobs2delete = []
		for i, it in enumerate(self.iterator):
			for tmp_i, tmp_j in enumerate(temp_jobs):
				new_job = self.duplicate_job(tmp_j, sufix_name = it)
				self.check_dependencies(new_job, it, temp_jobs)
				self.parse_iter(it, self.name, new_job)
				self.add_nested_iteration_relation(tmp_j, new_job)
				Batch.jobs_names.append(new_job.name)
				jobs.append(new_job)
				self.jobs.append(new_job)
				jobs2delete.append(tmp_i)
		self.delete_jobs(jobs2delete, temp_jobs) #Remove temporal jobs

	def duplicate_job(self, tmp_j, sufix_name = ''):
		new_job = copy.deepcopy(tmp_j)
		new_job.name = tmp_j.name+'_'+sufix_name
		new_job.attrib = copy.deepcopy(tmp_j.attrib)
		new_job.dependencies = copy.deepcopy(tmp_j.dependencies)
		new_job.initialization = copy.deepcopy(tmp_j.initialization)
		new_job.parameters = copy.deepcopy(tmp_j.parameters)
		return new_job

	def find_job_names(self, name):
		final_names = []
		intermediary_names = Batch.nested_iteration_relations.get(name)
		if intermediary_names != None:
			while len(intermediary_names) > 0:
				final_names = intermediary_names
				i_names = []
				for i_n in intermediary_name:
					query = Batch.nested_iteration_relations[i_n]
					if query != None: i_names.extend(query)
				if len(i_names > 0): 
					intermediary_names = i_names
				else: 
					break
		return final_names

	#tmp_j => job to set dependencies in iteration
	#iter => sufix of current iteration
	#jobs => array of jobs which has the job dependency	
	def check_dependencies(self, tmp_j, it, jobs):
		jobs_names = [ job.name for job in jobs ]
		deps = {}
		for i, dep in enumerate(tmp_j.dependencies): 
			if dep in jobs_names: deps[dep] = i

		for name, index in deps.items():
			dep = name+'_'+it
			tmp_j.initialization = re.sub(name+'\\)', dep+')', tmp_j.initialization)
			tmp_j.parameters = re.sub(name+'\\)', dep+')', tmp_j.parameters)
			tmp_j.dependencies[index] = dep 


	def parse_iter(self, it, name, job):
		job.parameters = self.set_iter(name, it, job.parameters)
		job.initialization = self.set_iter(name, it, job.initialization)

	def set_iter(self, name, it, string):
		string = re.sub(name+'\\(\\+\\)', it, string)
		return string

	def check_regex_dependencies(self):
		if self.regex_deps == 'tag':
			new_job_names = []
			for it in self.iterator:
				new_names = self.find_job_names(re.sub('\\)', '', it))
				new_job_names.extend(new_names)
			if len(new_job_names) > 0: self.iterator = [ nj + ')' for nj in new_job_names ]

	def add_nested_iteration_relation(self, tmp_j, new_job):
		query = Batch.nested_iteration_relations.get(tmp_j.name)
		if query == None:
			Batch.nested_iteration_relations[tmp_j.name] = [new_job.name]
		else:
			query.append(new_job.name)

	def delete_jobs(self, jobs2delete, job_array):
		jobs2delete = list(set(jobs2delete))
		for index in sorted(jobs2delete, reverse=True): job_array.pop(index)
