__author__ = 'LieuV'

from os.path import basename, splitext, join, isfile, dirname, abspath, exists
from os import makedirs, listdir, remove
import subprocess
import re, sys, numpy, shutil
import xml.etree.ElementTree as ET
from collections import OrderedDict
import json

def getDir(args, getAbs = False):
	if isfile(args[0]):
		args[0]= dirname(args[0])
	path = join(*args)
	if(getAbs):
		path = abspath(path)
	makedirs(path, exist_ok=True)
	return path

def getFilenameWithoutExtension(string, getBase = False):
	'''   In case of double extension   '''
	if getBase:
		string = basename(string)
	string = splitext(string)[0]
	i = 0
	while splitext(string)[-1] in [".gz", ".gzip", ".zip", ".bz", ".fasta", ".fastq", ".bam"]:
		string = splitext(string)[0]
	return string

def bam_to_fastq(bamfile, output):
	print("Convert bam to fastq...", bamfile)
	outfile = join(output, getFilenameWithoutExtension(bamfile, getBase= True)  +".fastq")
	subprocess.Popen("bamtools convert -format fastq -in " + bamfile + " -out " + outfile, shell = True)
	return outfile

def toLatex(text, newline = False):
	text = text.replace("#","\#")
	text = text.replace("_", "\_")
	text = text.replace("\t", " ")
	text = text.replace("%", "\%")

	if newline:
		text = text.replace("\n", " ")
	return text

def join_reads(read_list, outputDir, outputName):
	if len(read_list) ==0:
		return None
	elif len(read_list) ==1:
		return read_list[0]
	outfile = join(outputDir, getFilenameWithoutExtension(outputName) )
	outfile += splitext(read_list[0])[-1]
	#if read_list[0].endswith(".gz"):
	print("Concatenate files. This may take a while." ," ".join(read_list) + " > " + outfile)
	process = subprocess.call("cat " + " ".join(read_list) + " > " + outfile, shell=True)
	#else:
	#	print("Decompress and concatenate files. This may take a while.")
	#	process = subprocess.call("gzip --fast -c " + " ".join(read_list) + " > " + outfile, shell=True)
	if process != 0:
		sys.exit("Error with gzip")
	return outfile

def trimmomatic_report(process, type, write = True):
	logs =""
	survived = 0
	total = 0
	pTrimmedReads = 0
	nTrimmedReads = 0
	for line in iter(process.stdout.readline, b''):
		print(line.decode("utf-8"))
	for line in iter(process.stderr.readline, b''):
		line = line.decode("utf-8")
		print(line)
		if not (line.startswith("Using Long Clipping Sequence") or line.startswith("Trimmomatic")):
			if write:
				logs += toLatex(line) + "\n"
			if re.match("Input Read", line):
				if re.search("TrimmomaticPE", type):
					pattern = re.match("Input Read Pairs:\s+(?P<total>\d+)\s+"
									   "Both Surviving:\s+(?P<nSurvived>\d+) \((?P<pSurvived>\d+.\d+)%\)\s+"
									   "Forward Only Surviving:\s+(?P<forward>\d+) \(\d+.\d+%\)\s+"
									   "Reverse Only Surviving:\s+(?P<reverse>\d+) \(\d+.\d+%\).*", line)

					survived = float(pattern.group("pSurvived").replace(",", "."))
					total = int(pattern.group("total")) * 2
					nTrimmedReads = int(pattern.group("nSurvived"))*2 + int(pattern.group("forward")) + int(pattern.group("reverse"))
					pTrimmedReads = round(100 * nTrimmedReads/total, 2)
				else:
					pattern = re.match(".*Surviving: (?P<nSurvived>\d+) \((?P<survived>\d+.\d+)%\)", line)
					pattern2 = re.match("Input Reads: (?P<total>\d+)", line)

					survived = float(pattern.group("survived").replace(",","."))
					total = int(pattern2.group("total"))
					nTrimmedReads = int(pattern.group("nSurvived"))
					pTrimmedReads = survived

	return survived, total,  nTrimmedReads, pTrimmedReads, logs

def perc_slope(a,b,perc):
	if abs(numpy.diff([a,b]))>(b*float(perc)):
		return True
	return False

def optimize_trimming(fastqc_folder,perc=0.1):
	mytable=None
	filename = join(fastqc_folder, "fastqc_data.txt")
	with open(filename, "rb") as fastqcdata:
		table = []
		ifwrite = False
		for line in fastqcdata.readlines():
			line = line.decode()
			line = line.replace("\n", "")
			if line.startswith(">>END_MODULE") and ifwrite:
				try:
					dtype = {'names': table[0], 'formats': ['|S15', float, float, float, float]}
					mytable = numpy.asarray([tuple(x) for x in table[1:]], dtype=dtype)
					break
				except:
					pass
			elif line.startswith("Sequence length"):
				seq_length = int(line.split("\t")[-1].split("-")[-1])
			elif line.startswith(">>Per base sequence content"):
				ifwrite = True
				temp = line.split("\t")

				if temp[1].lower() == "pass":
					return False, 0, seq_length  #break
			elif ifwrite:
				table.append(line.split("\t"))

	headcrop = 0
	tailcrop = len(mytable["A"])
	trim_bool = False
	column = numpy.ma.array(mytable)
	for i in range(-4, int(round(len(mytable["A"]) / 3, 0)), 1):
		for nucl in ["A", "C", "G", "T"]:
			column[nucl].mask[max(i, 0):i + 5] = True
			if headcrop >0:
				column[nucl].mask[:headcrop] = True
			if tailcrop < len(mytable["A"]):
				column[nucl].mask[tailcrop:] = True

			# check heacrop
			if (perc_slope(numpy.mean(mytable[nucl][max(i, 0):i + 5]), numpy.mean(column[nucl]), perc=perc)) & (headcrop < (i + 5)):
				headcrop = i + 5
				trim_bool = True
			elif headcrop < i:
				column[nucl].mask[max(i, 0):i + 5] = False

			# now crop from the end
			column[nucl].mask[-(i + 5):(min(len(mytable[nucl]), len(mytable[nucl]) - i))] = True
			if (perc_slope(numpy.mean(mytable[nucl][-(i + 6): (min(len(mytable[nucl]) - 1, len(mytable[nucl]) - 1 - i))]),numpy.mean(column[nucl]), perc=perc)) & (tailcrop > len(mytable[nucl]) - (i + 5)):
				# if perc_slope(numpy.var(mytable[nucl][-(i+6): (min(len(mytable["A"])-1, len(mytable[nucl])-1-i))]), numpy.mean(numpy.var(column)), perc=perc):
				#print("TAILCROP", tailcrop)
				tailcrop = len(mytable[nucl]) - (i + 5)
				trim_bool = True
			else:
				column[nucl].mask[-(i + 5): (min(len(mytable["A"]) - 1, len(mytable[nucl]) - 1 - i))] = False

	return trim_bool, headcrop, tailcrop-headcrop

def str_to_json(report_str):
	closings = []
	json_file = ""
	report_str = report_str.split("\n")
	pre_depth = -1
	for row in report_str:
		nNode = 0
		if row != "":
			row = row.split("\t")
			if row[0].strip() != "0.00" and row[5] != "unclassified":
				temp = len(row[5])
				name = row[5].strip(" ")
				curr_depth = int((temp - len(name)) / 2)

				if pre_depth < curr_depth and curr_depth > 0:
					for i in range(curr_depth - pre_depth):
						if json_file.endswith('"children":[ '):
							if json_file.endswith("}") or json_file.endswith("]"):
								json_file += ","
							json_file += '{"children":[ '
							closings.append("]")
							closings.append("}")
						else:
							json_file += ',"children":[ '
							closings.append("]")
				elif pre_depth > curr_depth:
					while pre_depth >= curr_depth:
						close_with = closings.pop()
						json_file += close_with  # end node
						if close_with == "]":
							close_with = closings.pop()
							json_file += close_with  # end node
						pre_depth -= 1
				elif pre_depth != -1:
					json_file += closings.pop() + ','
				if json_file.endswith("}")or json_file.endswith("]"):
					json_file+=","
				json_file += '{"name":"%s",' % name
				json_file += '"perc":%s,' % row[0]
				json_file += '"taxa":"%s"' % row[3]
				closings.append("}")
				nNode += 1
				pre_depth = curr_depth
	while len(closings) > 0:
		json_file += closings.pop()  # end node
	return json_file

def delete_files(path, extended="", pattern="_fastqc"):
	for f in listdir(join(path, extended)):
		filename =join(path,extended,f)
		if f.endswith(pattern):
			if isfile(filename):
				remove(filename)
			else:
				shutil.rmtree(filename)

'''
def get_distribution(lines, key):
	#creates CSV used for boxplotting
	output = []
	for line in lines:
		line = line.strip()
		fields = line.split('\t')
		if key == "Sequence_Length_Distribution":
			l = fields[0].split('-')
			try:
				mean = (int(l[0]) + int(l[1])) / 2
			except:
				mean = int(l[0])
			output.extend([str(mean)] * int(float(fields[1])))

			# sequence quality or sequence GC content
		if key == 'Per_sequence_GC_content' or key == "Per_sequence_quality_scores":
			v = fields[0]  # quality or gc [%]
			output.extend(int(v) * [int(float(fields[1]))])
	return output
'''
def createCSV(dataset, lines, key, path, trimmed):
	'''creates CSV used for boxplotting'''
	#output = []
	splitted_ds=dataset.split("_")

	csvhandle = open(path + '/' + key + ".csv", 'a')
	for line in lines:
		line = line.strip()
		fields = line.split('\t')
		#Sequence Length Distribution
		if key ==  "Sequence_Length_Distribution":
			l = fields[0].split('-')
			try:
				value = (int(l[0]) + int(l[1])) / 2
			except:
				value = int(l[0])
			count = int(float(fields[1]))

		#sequence quality or sequence GC content
		if key == 'Per_sequence_quality_scores' or key == 'Per_sequence_GC_content':
			value = fields[0]  #quality or gc [%]
			count = int(float(fields[1]))
		try:
			if splitted_ds[-2] == "R1" or splitted_ds[-2] == "R2":
				# print(splitted_ds)
				csvhandle.write(",".join(
					["_".join(splitted_ds[:-3]), splitted_ds[-2], splitted_ds[-3], trimmed, str(value),
					 str(count)]) + "\n")
			else:
				csvhandle.write(",".join(
					["_".join(splitted_ds[:-4]), splitted_ds[-3], splitted_ds[-4], trimmed, str(value),
					 str(count)]) + "\n")
		except:
			csvhandle.write(",".join([dataset.replace("_fastqc",""), "NA", "NA" , trimmed, str(value), str(count)]) + "\n")
	csvhandle.close()

def write_fastqc_summary(fastqcdir, output):
	store = False
	data = []
	fastqcdatafile = open(join(fastqcdir,"fastqc_data.txt"), "r")
	for line in iter(fastqcdatafile):
		if line.startswith("#"):
			pass  # print("comment")
		# continue
		elif line.startswith("Total Sequences"):
			total_seq = int(line.split("\t")[-1])
		elif line.startswith('>>END_MODULE'):
			if len(data) > 0:
				name1 = basename(fastqcdir)
				if "Trimmed" in fastqcdir:
					# name1 = "Trimmed " + name1
					trimmed = "Trimmed"
				else:
					trimmed = "Raw"
				createCSV(name1, data[2:], key,  output, trimmed)
				data = []
				store = False
		elif line.startswith('>>'):
			key = line.split("\t")[0]
			key = key.replace(">>","")
			key = key.replace(" ","_")
			if key in ['Sequence_Length_Distribution', 'Per_sequence_quality_scores','Per_sequence_GC_content']:
				store = True
		if store:
			data.extend([line])
	fastqcdatafile.close()
	return total_seq

def write_SAV(folder, output):
	print("Plot SAV images...")
	process = subprocess.Popen(" ".join(["Rscript --vanilla ", join(dirname(__file__), "sav.R"),
										folder, join(output,"img")]),
							  stderr=subprocess.PIPE, stdout=subprocess.PIPE, shell=True)
	for line in iter(process.stderr.readline, b''):
		print(line)
	tree = ET.parse(join(folder,'RunInfo.xml'))
	root = tree.getroot()
	runinfo = OrderedDict()

	for child in root.iter("Run"):
		for grand in child.getchildren():

			if len(grand.attrib) != 0:
				runinfo[grand.tag] = grand.attrib

			if len(grand.getchildren()) !=0:
				for grandgrand in grand:
					if len(grandgrand.attrib) != 0:
						try:
							runinfo[grand.tag].append(grandgrand.attrib)
						except:
							runinfo[grand.tag] = [grandgrand.attrib]

					if grandgrand.text is not None:
						if grandgrand.text.strip() != "":
							if grand.tag not in runinfo.keys():
								runinfo[grand.tag]={}
							try:
								runinfo[grand.tag][grandgrand.tag].append(grandgrand.text)
							except:
								runinfo[grand.tag][grandgrand.tag] = [grandgrand.text]

			if grand.text is not None:
				if grand.text.strip()!="":
					if len(grand.attrib) != 0:
						runinfo[grand.tag][grand.tag] = grand.text
					else:
						runinfo[grand.tag] = grand.text

	if exists(join(folder, 'RunParameters.xml')):
		runparameter_xml = ET.parse(join(folder, 'RunParameters.xml'))
	elif exists(join(folder, 'runParameters.xml')):
		runparameter_xml = ET.parse(join(folder, 'runParameters.xml'))
	else:
		sys.exit("No RunParameter.xml found.")

	runparameters_unique = ["ExperimentName", "FPGAVersion", "CPLDVersion", "MCSVersion","RTAVersion", "CameraFirmware", "CameraDriver", "ClusteringChoice" ,"RapidRunChemistry", "RunMode", "ApplicationName", "ApplicationVersion"]
	runparameters_path = ["Setup/ReagentKits/Sbs/SbsReagentKit/ID","Setup/ReagentKits/Index/ReagentKit/ID","Setup/Sbs","Setup/Rehyb", "PR2BottleRFIDTag/PartNumber","FlowcellRFIDTag/PartNumber","ReagentKitRFIDTag/PartNumber"]
	root = runparameter_xml.getroot()
	runparameter = OrderedDict()
	for label in runparameters_unique:
		for i in root.iter(label):
			runparameter[label] =  i.text
			#print(label, ": ", i.text)

	for label_path in runparameters_path:
		label_path = label_path.split("/")
		element = root.find(label_path[0])
		for label in label_path[1:]:
			if element is None:
				break
			element = element.find(label)
		if element is not None:
			runparameter["/".join(label_path)] = element.text
			#print(label_path, element.text)
		else:
			print(label_path, " does not exist.")
	runparameter.update(runinfo)
	json.dump({"tables":runparameter}, open(join(output, "sav.json"),"w"))