#!/usr/bin/env python3
##########################################################
## Jose F. Sanchez					##
## Copyright (C) 2019-2020 Lauro Sumoy Lab, IGTP, Spain	##
##########################################################
"""
Shared functions used along ``BacterialTyper`` pipeline.
With different purposes:

	- Print time stamps

	- Create system calls

	- Manage/list/arrange files/folders

	- Aesthetics

	- Manage fasta files

	- Other miscellaneous functions
"""
## useful imports
import time
import io
import os
import re
import subprocess
import sys
import wget
import json
from datetime import datetime
from Bio import SeqIO
from termcolor import colored
from filehash import FileHash
import pandas as pd
import patoolib ## to extract
import shutil

## import my modules
from BacterialTyper.config import set_config
from BacterialTyper.other_tools import tools

########################################################################
######## 					TIME								######## 					
########################################################################

###############   
def print_time ():
	"""Prints time stamp in human readable format: month/day/year, hour:minute:seconds."""
	now = datetime.now()
	date_time = now.strftime("%m/%d/%Y, %H:%M:%S")
	print ('\t' + date_time + '\n')

###############   
def gettime (start_time):
	"""Obtains time stamp in human readable format: hour:minute:seconds from a time.time() format timestamp."""
	total_sec = time.time() - start_time
	m, s = divmod(int(total_sec), 60)
	h, m = divmod(m, 60)
	return h, m, s

###############	
def timestamp (start_time_partial):
	"""Prints a stamp of the time spent for a process in human readable format: hour:minute:seconds.
	Returns time in format time.time().
	"""
	h,m,s = gettime(start_time_partial)
	print_sepLine("-", 25, False)
	print ('(Time spent: %i h %i min %i s)' %(int(h), int(m), int(s)))
	print_sepLine("-", 25, False)
	return time.time()

###############	
def print_time_stamp (out):
	"""Prints out timestamp in a file provided. Format: time.time()"""
	timefile = open(out, 'w')    
	string2write = str(time.time())
	timefile.write(string2write)
	return()

###############	
def read_time_stamp (out):
	"""Reads timestamp from a file provided. Format: time.time()"""
	st_hd = open(out, 'r')
	st = st_hd.read()
	st_hd.close()
	stamp = datetime.fromtimestamp(float(st)).strftime('%Y-%m-%d %H:%M:%S')
	return(stamp)

###############	
def get_diff_time(stamp):
	"""Obtains the time spent for a process in days given a stamp in time.time() format.
	Returns days passed since.
	"""

	time_today = time.time()
	elapsed = time_today - float(time_today)
	days_passed = int((elapsed/3600)/24)
	return(days_passed)

###############    
def create_human_timestamp():
	"""Generates human timestamp for the date of day in format (yearmonthday): e.g. 20191011"""
	now = datetime.now()
	timeprint = now.strftime("%Y%m%d")
	return timeprint
###############

############################################################################
######## 					FILES/FOLDERS							######## 					
############################################################################

###############
def is_non_zero_file(fpath):  
	# https://stackoverflow.com/a/15924160
	"""Returns TRUE/FALSE if file exists and non zero"""
	return os.path.isfile(fpath) and os.path.getsize(fpath) > 0

###############
def outdir_project(outdir, project_mode, pd_samples, mode):
	"""
	"""
	# Group dataframe by sample name
	sample_frame = pd_samples.groupby(["new_name"])

	dict_outdir = {}	
	for name, cluster in sample_frame:
		if (project_mode):
			#print ("Create subdir for every sample: ", mode)
			sample_dir = create_subfolder('data', outdir)		

			## create sample
			sample_name_dir = create_subfolder(name, sample_dir)		

			## create subdir sub sample
			mode_name_dir = create_subfolder(mode, sample_name_dir)		
			dict_outdir[name] = mode_name_dir

		else:
			#print ("All samples share same folder")
			sample_name_dir = create_subfolder(name, outdir)		
			dict_outdir[name] = sample_name_dir

	return (dict_outdir)

###############
def outdir_subproject(outdir, pd_samples, mode):
	## we assume we want to create within a project dir a subdir
	# Group dataframe by sample name
	sample_frame = pd_samples.groupby(["name"])
	dict_outdir = {}	
	for name, cluster in sample_frame:
		mode_name_dir = create_subfolder(mode,os.path.join(outdir, name))		
		dict_outdir[name] = mode_name_dir

	return (dict_outdir)

###############
def create_subfolder (name, path):
	"""Create a subfolder named 'name' in directory 'path'. Returns path created."""
	## create subfolder  ##	
	subfolder_path = path + "/" + name
	access_rights = 0o755

	# define the access rights
	try:
		os.mkdir(subfolder_path, access_rights)
	except OSError:  
	   	#print ("\tDirectory %s already exists" % subfolder_path)
		return subfolder_path
	else:  
		print (colored("Successfully created the directory %s " % subfolder_path, 'yellow'))

	return subfolder_path

###############  
def create_folder (path):
	"""Create a folder directory 'path'. Returns path created."""

	## create subfolder  ##	
	access_rights = 0o755

	# define the access rights
	try:
		os.mkdir(path, access_rights)
	except OSError:  
		#print ("\tDirectory %s already exists" %path)
		return path
	else:  
		print (colored("Successfully created the directory %s " %path, 'yellow'))

	return path

############### 
def get_symbolic_link (sample_list, directory):
	"""Creates symbolic links, using system call, for list of files given in directory provided"""
	for samplex in sample_list:
		cmd = 'ln -s %s %s' %(samplex, directory)
		system_call(cmd, returned=False)

	files2return = os.listdir(directory)
	return files2return

############### 
def get_symbolic_link_file (file2link, newfile):
	"""Creates symbolic link for a file into a new name file"""
	cmd = 'ln -s %s %s' %(file2link, newfile)
	system_call(cmd, returned=False)

#################
def get_fullpath_list(dir_given):
	"""Retrieve full absolute path for the files within a directory specified.

	:param dir_given: Directory to retrieve files
	:type dir_given: string

	:returns: List of absolute path files.
	"""
	return_path = []
	for root, dirs, files in os.walk(dir_given):
		for f in files:
			return_path.append(os.path.join(root,f))
	return return_path

#################
def printList2file(fileGiven, listGiven):
	"""Prints list given in the output file provided. One item per row."""
	file_hd = open(fileGiven, 'w')
	file_hd.write("\n".join(listGiven))
	file_hd.close()  

#################
def readList_fromFile(fileGiven):
	"""Reads list from the input file provided. One item per row."""
	# open file and read content into list
	lineList = [line.rstrip('\n') for line in open(fileGiven)]
	return (lineList)

######
def retrieve_matching_files(folder, string):
	"""Lists folder path provided and given a string to search, returns all files ending with the given string"""
	my_all_list = get_fullpath_list(folder)
	matching = [s for s in my_all_list if s.endswith(string)]
	return (matching)

#####
def file2dictionary(file2read, split_char):
	"""Read file and generate a dictionary"""
	d = {}
	with open(file2read) as f:
		for line in f:
			line = line.rstrip('\n')
			(key, val) = line.split(split_char)
			d[key] = val

	return(d)

def file2dataframe(file2read, names):
	## TODO: check if duplicated with get_data function
	"""Read csv file into pandas dataframe"""
	d = pd.read_csv(file2read, comment="#", names=names)
	return(d)


def chmod_rights(file, access_rights):
	"""Changes access permission of a file:
	
	Read additional information in:
	https://docs.python.org/3/library/os.html#os.chmod
	https://www.geeksforgeeks.org/python-os-chmod-method/
	"""
	access = os.chmod(file, access_rights, follow_symlinks=True)


########################################################################
######## 					system call							######## 					
########################################################################

###############
def system_call(cmd, returned=False, message=True):
	"""Generates system call using subprocess.check_output"""
	## call system
	## send command
	if (message):
		print (colored("[** System: %s **]" % cmd, 'magenta'))

	try:
		out = subprocess.check_output(cmd, shell = True)
		if (returned):
			return (out)
		return ('OK')
	except subprocess.CalledProcessError as err:
		if (returned):
			return (err.output)
		if (message):
			print (colored("** ERROR **", 'red'))
			print (colored(err.output, 'red'))
			print (colored("** ERROR **", 'red'))
		
		return ('FAIL')

###############
def wget_download(url, path):
	"""
	Downloads file in the given path.
	
	It uses python module wget to download given ``url`` in the ``path`` provided.
	
	:param url: File to download from http/ftp site.
	:param path: Absolute path to the folder where to save the downloaded file.
	:type url: string
	:type path: string  
	
	:returns: Print messages and generates download
	"""
	
	print ('\t+ Downloading: ', url)
	wget.download(url, path)
	print ('\n')

###############
def check_md5sum(string, File):
	md5hasher = FileHash('md5')
	md5_file = md5hasher.hash_file(File)
	
	if (md5_file == string):
		#print (md5_file + '==' + string)
		return (True)
	else:
		return (False)

###############
def extract(fileGiven, out, remove=True):
	"""
	Extracts archived file
	
	This function extracts the file given in the ``out`` path provided.
	It uses the ``patoolib`` that is able to identify the type of file 
	and compression algorithm to use in each case.
	
	It also removes the compressed file using ``os`` module.
	
	:param fileGiven: Archived file to extract.
	:param out: Output name and absolute path for the extracted archived.
	:param remove: True/False for removing compressed file extracted
	
	:type fileGiven: string
	:type out: string
	:type remove: boolean
	
	"""
	## extract using patoolib
	patoolib.extract_archive(fileGiven, outdir=out, verbosity=0)
	
	if (remove):
		## remove compress file
		print ("Remove compress file...")
		os.remove(fileGiven)
		print ("\n")
	

####################################################################
######## 					BLAST							######## 					
####################################################################

###############
def makeblastdb(DBname, fasta):
	## generate blastdb for genome
	makeblastDBexe = set_config.get_exe('makeblastdb')
	
	if (os.path.isfile(DBname + '.nhr')):
		print ("+ BLAST database is already generated...")
	else:
		cmd_makeblast = "%s -in %s -input_type fasta -dbtype %s -out %s" %(makeblastDBexe, fasta, 'nucl', DBname)
		code = system_call(cmd_makeblast)

		if (code == 'FAIL'):
			print (colored('****ERROR: Some error happened during the makeblastDB command', 'red'))
			print (cmd_makeblast)
			exit()
	

###############	
def blastn(outFile, DBname, fasta, threads):
	# blastn plasmids vs contigs
	blastnexe = set_config.get_exe('blastn')
	cmd_blastn = "%s -db %s -query %s -out %s -evalue 1e-20 -outfmt \'6 std qlen slen\' -num_threads %s" %(blastnexe, DBname, fasta, outFile, threads )
	codeBlastn = system_call(cmd_blastn)
	
	if (codeBlastn == 'FAIL'):
		print (colored('****ERROR: Some error happened during the blastn command', 'red'))
		print (cmd_blastn)
		exit()
	
########################################################################
######## 					FASTA files							######## 					
########################################################################

###############
def concat_fasta(dirFasta, Fasta):
	print ("+ Concatenating all information into one file...")	
	cmd = 'cat ' + dirFasta + '/*fna > ' + Fasta
	return(system_call(cmd))

###############
def subset_fasta(ident, fasta, out):
	output_FASTA = open(out, 'w')	
	for record in SeqIO.parse(fasta, "fasta"):
		all_id = record.description
		species_search = re.search(r"%s" % ident, all_id)
		if species_search:
			head = ">" + all_id + "\n"
			output_FASTA.write(head)
			output_FASTA.write(str(record.seq))
			output_FASTA.write("\n")
	
	output_FASTA.close()

###############
def rename_fasta_seqs(fasta_file, name, new_fasta):
	"""Rename fasta sequences provided in file :file:`fasta_file` using id :file:`name`. Save results in file :file:`new_fasta` provided.
	
	Check for id character lenght do not exceed 37 characters as it might be a limitiation in further annotation and subsequent analysis. Read Prokka_ issue for further details: https://github.com/tseemann/prokka/issues/337.
	
	:param fasta_file: Absolute path to fasta file.
	:type fasta_file: string
	:param name: String to add every fasta sequence header.
	:type name: string
	:param new_fasta: Name for the new fasta file (Absolute path).
	:type new_fasta: string
	:return: Path to tabular delimited file containing conversion from all to new id for each sequence.
	:warnings: Returns FAIL if name is >37 characters.
	
	.. include:: ../../links.inc	 	
	"""

	output_FASTA = open(new_fasta, 'w')
	id_conversion = open(new_fasta + "_conversionID.txt", 'w')
	
	if len(name) > 37:
		print (colored("** ERROR **", 'red'))
		print (colored("BacterialTyper.functions.rename_fasta_seqs():: name id is > 37 characters.", 'red'))
		print (colored("** ERROR **", 'red'))
		return ('FAIL')
	
	counter_seqs = 0
	for record in SeqIO.parse(fasta_file, "fasta"):
		old_id = record.description
		counter_seqs += 1
		new_id = name + "_" + str(counter_seqs)
		head = ">" + new_id + "\n"
		output_FASTA.write(head)
		output_FASTA.write(str(record.seq))
		output_FASTA.write("\n")
		
		id_conversion.write(old_id + "\t" + new_id)
		id_conversion.write("\n")
	
	output_FASTA.close()
	id_conversion.close()
		
	return (new_fasta + "_conversionID.txt")
	

############################################################################
######## 					AESTHETICS								######## 					
############################################################################
def pipeline_header():
	"""
	Prints a common header for the pipeline including name, author, copyright and year.	    
	"""
	print ("\n")
	print_sepLine("#", 70, False)
	print('#', '{: ^66}'.format("BacterialTyper pipeline"), '#')
	print('#', '{: ^66}'.format("Jose F. Sanchez, Cristina Prat & Lauro Sumoy"), '#')
	print('#', '{: ^66}'.format("Copyright (C) 2019-2020 Lauro Sumoy Lab, IGTP, Spain"), '#')
	print_sepLine("#", 70, False)

###############
def print_sepLine(char, num, color):
	string = char * num
	if (color):
		print (colored(string, color))
	else:
		print (string)

###############
def boxymcboxface(message):
	## this function is from ARIBA (https://github.com/sanger-pathogens/ariba)
	## give credit to them appropiately
   	#print('-' * 79)
	print ('\n')
	print('|', '=' * 50, '|', sep='')
	print('|', '{: ^48}'.format(message), '|')
	print('|', '=' * 50, '|', sep='')
	print ('\n')
    #print('-' * 79)

###############
def progbar(curr, total, full_progbar):
	frac = (curr/total)
	filled_progbar = round(frac*full_progbar)
	print ('\r', '#'*filled_progbar + '-'*(full_progbar-filled_progbar), '[{:>7.2%}]'.format(frac), end='')
	sys.stdout.flush()


########################################################################
######## 					Miscellaneous						######## 					
########################################################################

###############
def get_data(ID_file, SEP, options):	
	if options == 'index_col=0':
		data = pd.read_csv(ID_file, sep=SEP, index_col=0)
		return(data)
	else:
		data = pd.read_csv(ID_file, sep=SEP)
		return(data)
	## fix for another example if any
	
###############
def get_number_lines(input_file):	
	with open(input_file) as foo:
		lines = len(foo.readlines())
	foo.close()
	return (lines)

###############
def get_info_file(input_file):	
	with open(input_file) as f:
	    lines = f.read().splitlines() 
	f.close()
	return (lines)
  	    
#################
def optimize_threads(total, samples):
	cpu = int(int(total)/int(samples))
	
	if (cpu==0): ## 5 availabe cpus for 10 samples == 0 cpu
		cpu = 1
	
	return(cpu)

#################
def parse_sublist(lst, ind): 
	## extract elemnts of the list
	## Original Code: https://www.geeksforgeeks.org/python-get-first-element-of-each-sublist/
	return [item[ind] for item in lst]

##################
def decode(x):
	"""
	Python decode string method

	It converts bytes to string.

	:param x: String of text to decode
	:type x: string
	:returns: Text decoded

	.. attention:: Be aware of Copyright

		The code implemented here was retrieved and modified from ARIBA (https://github.com/sanger-pathogens/ariba)

		Give them credit accordingly.
	"""
	try:
		s = x.decode()
	except:
		return x

	return s

###################
def collect_info_run(out_folder, module, userinput, runInfo):
	"""
	Prints information of the job run in json format
	
	:param out_folder: Absolute path to folder to dump results
	:param module: Module or script name executed
	:param userinput: Contains information about the input parameters
	:param runInfo: Contains configuration detail information to dump
	
	:type out_folder: string 
	:type module: string 
	:type userinput: dict
	:type runInfo: dict
	
	:returns: Prints information in json format in the output folder provided.
	
	Exampple of infromation to include:
	
	userinput = {"filename":infiles, 
				"database":path_to_database,
				"file_format":file_format}
				
	runInfo = { "module":module, 
				"analysis":example,
				"date":date, 
				"time":time}
	
	Original idea extracted from https://bitbucket.org/genomicepidemiology/kmerfinder/src/master/kmerfinder.py
	
	"""
	
	# Collect all data of run-analysis
	data = {service:{}}
	
	data[service]["user_input"] = userinput
	data[service]["run_info"] = runInfo
	
	# Save json output
	result_json_file = out_folder + "/" + module + ".json" 
	with open(result_json_file, "w") as outfile:  
	   json.dump(data, outfile)
	
def print_all_pandaDF(pd_df):
	pd.set_option('display.max_colwidth', None)
	pd.set_option('display.max_columns', None)
	print (pd_df)
