"""
FUGAsseM utility: prepare annotation and evidence for as FUGAsseM inputs
A collection of tasks for preparing annotation files that will be used as FUGAsseM inputs

Copyright (c) 2021 Harvard School of Public Health

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import sys
import os
from os import path
import subprocess
import itertools
import re

from anadama2.tracked import TrackedExecutable, TrackedDirectory
from anadama2 import Workflow

# import the utilities functions and config settings from MetaWIBELE
try:
	from fugassem import utilities, config
except ImportError:
	sys.exit("CRITICAL ERROR: Unable to find the FUGAsseM python package." +
		         " Please check your install.")


VERSION = config.version

def parse_cli_arguments():
	'''
	Parses any command-line arguments passed into the workflow.
	the version number will appear when running this script with the "--version" option
	the description will appear when running this script with the "--help" option
	create a workflow instance, providing the version number and description
	'''

	working_dir = os.getcwd()
	workflow = Workflow(version=VERSION, description="FUGAsseM utility for preparing annotations",
	                    remove_options=["output"])

	# add the custom arguments to the workflow
	workflow.add_argument("func-type",
	                      desc = "type of function, [ Default: GO ]",
						  choices = ["GO", "BP", "MF", "CC",
							"UniRef90_GO", "UniRef90_GO_BP", "UniRef90_GO_CC", "UniRef90_GO_MF", "UniRef90_COG", "UniRef90_eggNOG",
                			"UniRef90_KEGG-KOs", "InterProScan_PfamDomain", "Denovo_transmembrane", "Denovo_signaling", "DOMINE_interaction"],
	                      default = "GO")
	workflow.add_argument("pfam",
	                      desc = "if specified, will extract pfam annotation of that type, [ Default: None ]",
	                      default = None)
	workflow.add_argument("ddi",
	                      desc = "if specified, will extract DDI annotation of that type, [ Default: DOMINE_interaction ]",
	                      default = "DOMINE_interaction")
	workflow.add_argument("contig",
	                      desc = "if specified, will extract contig annotation",
	                      action = "store_true")
	workflow.add_argument("clust-file",
	                      desc = "input clustering file in extended-fasta format for protein families, which can be generated by MetaWIBELE [ Default: None ]",
	                      default = None)
	workflow.add_argument("gene-info",
	                      desc = "input gene info file including the contig info of each member in protein families [ Default: None ]",
	                      default = None)
	workflow.add_argument("coann-list",
	                      desc = "specify which typs of data will converted into coann pairs as matrix evidence (pfam|ddi|contig), use comma to seprate [ Default: None ]",
	                      default = None)
	workflow.add_argument("homology",
	                      desc = "If specified, will extract co-homology annotation",
	                      action = "store_true")
	workflow.add_argument("homology-ann",
	                      desc = "homology-based annotation file [ Default: None ]",
	                      default = None)
	workflow.add_argument("basename",
	                      desc = "specify the basename for output files [ Default: fugassem ]",
	                      default = "fugassem")
	workflow.add_argument("threads",
	                      desc = "number of threads/cores for each task to use",
	                      default = 1)
	workflow.add_argument('memory',
	                      desc = 'The amount of memory to use for each fugassem job. Provided in MB',
	                      default = '10240')
	workflow.add_argument('time',
	                      desc = 'The amount of time to use for each fugassem job. Provided in minute',
	                      default = '600')
	workflow.add_argument("output",
	                      desc = "provide an output folder which the workflow database and log is written. "
	                           "By default, that be written to the anadama2 folder of users' working directory",
	                      default = working_dir)

	return workflow



def prepare_annotation (ann_file, ann_type, ann_map_file,
                output_folder, extracted_ann_file, extracted_ann_file2,
                workflow, threads, time_equation, mem_equation):

	"""
	Extract annotation for specific type

	Args:
		ann_file: MetaWIBELE finalized characterization file.
		ann_type: annotation type of interest.
		ann_map_file: function name mapping file [Default: None].
		output_folder (string): The path of the output folder.
		extracted_ann_file: output file name for specific annotation.
		workflow (anadama2.workflow): An instance of the workflow class.
		threads (int): The number of threads/cores for clustering to use.
		time_equation (int): required number of hours defined in the workflow.
		mem_equation (int): required number of GB defined in the workflow.

	Requires:
		MetaWIBELE characterization file.

	Returns:
		string: Annotation file

	Example:
		from anadama2 import Workflow
		from fugassem.tools.generate_annotation_input import prepare_annotation

		# create an anadama2 workflow instance
		workflow=Workflow()

		# add prepare_function tasks
		specific_ann_file = prepare_annotation (
						ann_file,
						ann_type,
						ann_map_file,
						output_folder,
						extracted_ann_file1,
						extracted_ann_file2,
						workflow,
                        args.threads,
                        args.time_equation,
                        args.mem_equation)
		# run the workflow
		workflow.go()
	"""

	config.logger.info("###### Start prepare_annotation module #####")

	# prep I/O files
	main_folder = output_folder
	if not os.path.isdir(main_folder):
		os.system("mkdir -p " + main_folder)
	extracted_ann_file = os.path.join(main_folder, os.path.basename(extracted_ann_file))
	extracted_log = extracted_ann_file + ".log"

	workflow.add_task(
		"fugassem_prepare_annotation -a [depends[0]] -m [args[0]] -t [args[1]] -o [targets[0]] > [args[2]] 2>&1",
		depends = [ann_file, TrackedExecutable("fugassem_prepare_annotation")],
		targets = [extracted_ann_file, extracted_ann_file2],
		args = [ann_map_file, ann_type, extracted_log],
		cores = 1,
		time = time_equation,
		mem = mem_equation,
		name = "fugassem_prepare_annotation")

	return extracted_ann_file


def prepare_coann_pairs (ann_file,
                output_folder, extracted_ann_file,
                workflow, threads, time_equation, mem_equation):

	"""
	Prepare co-annotation pairs for specific type

	Args:
		ann_file: Raw annotation file.
		output_folder (string): The path of the output folder.
		extracted_ann_file: output file name for co-annotation pairs.
		workflow (anadama2.workflow): An instance of the workflow class.
		threads (int): The number of threads/cores for clustering to use.
		time_equation (int): required number of hours defined in the workflow.
		mem_equation (int): required number of GB defined in the workflow.

	Requires:
		Raw annotation file

	Returns:
		string: the file of co-annotation file.

	Example:
		from anadama2 import Workflow
		from fugassem.tools.generate_annotation_input import prepare_annotation

		# create an anadama2 workflow instance
		workflow=Workflow()

		# add prepare_function tasks
		coann_file = prepare_coann_pairs (
						ann_file,
						output_folder,
						extracted_ann_file,
						workflow,
                        args.threads,
                        args.time_equation,
                        args.mem_equation)
		# run the workflow
		workflow.go()
	"""

	config.logger.info("###### Start prepare_coann_pairs module #####")

	# prep I/O files
	main_folder = output_folder
	if not os.path.isdir(main_folder):
		os.system("mkdir -p " + main_folder)
	extracted_ann_file = os.path.join(main_folder, os.path.basename(extracted_ann_file))
	extracted_log = extracted_ann_file + ".log"

	workflow.add_task(
		"fugassem_prepare_coann_pairs -i [depends[0]] -o [args[0]] > [args[1]] 2>&1",
		depends = [ann_file, TrackedExecutable("fugassem_prepare_coann_pairs")],
		targets = [extracted_ann_file + ".gz"],
		args = [extracted_ann_file, extracted_log],
		cores = 1,
		time = time_equation,
		mem = mem_equation,
		name = "fugassem_prepare_coann_pairs")

	return extracted_ann_file


def prepare_contig (clust_file, gene_info_file,
                output_folder, contig_ann_file, contig_ann_file2,
                workflow, threads, time_equation, mem_equation):

	"""
	Extract contig annotation for gene families

	Args:
		clust_file: clustering file in extended-fasta format for protein families, which can be generated by MetaWIBELE
		gene_info_file: gene info file including the contig info of each member in protein families.
		output_folder (string): The path of the output folder.
		extracted_ann_file: output file name for specific annotation.
		workflow (anadama2.workflow): An instance of the workflow class.
		threads (int): The number of threads/cores for clustering to use.
		time_equation (int): required number of hours defined in the workflow.
		mem_equation (int): required number of GB defined in the workflow.

	Requires:
		annotation file

	Returns:
		string: the file of MetaWIBELE characterization file.

	Example:
		from anadama2 import Workflow
		from fugassem.tools.generate_annotation_input import prepare_contig

		# create an anadama2 workflow instance
		workflow=Workflow()

		# add prepare_function tasks
		contig_ann_file = prepare_contig (
						clust_file,
						gene_info_file,
						output_folder,
						contig_ann_file,
						workflow,
                        args.threads,
                        args.time_equation,
                        args.mem_equation)
		# run the workflow
		workflow.go()
	"""

	config.logger.info("###### Start prepare_contig module #####")

	# prep I/O files
	main_folder = output_folder
	if not os.path.isdir(main_folder):
		os.system("mkdir -p " + main_folder)
	contig_ann_file = os.path.join(main_folder, os.path.basename(contig_ann_file))
	extracted_log = contig_ann_file + ".log"

	workflow.add_task(
		"fugassem_prepare_contig -c [depends[0]] -m [depends[1]] -o [targets[0]] > [args[0]] 2>&1",
		depends = [clust_file, gene_info_file, TrackedExecutable("fugassem_prepare_contig")],
		targets = [contig_ann_file, contig_ann_file2],
		args = [extracted_log],
		cores = 1,
		time = time_equation,
		mem = mem_equation,
		name = "fugassem_prepare_annotation")

	return contig_ann_file


def prepare_co_unit (ann_file, ann_type,
                      unit_file, func_file,
                      output_folder, extracted_ann_file,
                      workflow, threads, time_equation, mem_equation, func_type):

	"""
	Extract co-unit info, e.g. co-homology

	Args:
		ann_file: MetaWIBELE finalized characterization file.
		ann_type: annotation type of interest.
		unit_file: unit annotation file.
		func_file: function annotation file.
		output_folder (string): The path of the output folder.
		extracted_ann_file: output file name for specific annotation.
		workflow (anadama2.workflow): An instance of the workflow class.
		threads (int): The number of threads/cores for clustering to use.
		time_equation (int): required number of hours defined in the workflow.
		mem_equation (int): required number of GB defined in the workflow.
		func_type: function type 

	Requires:
		annotation file
		unit file
		function file

	Returns:
		string: the file of MetaWIBELE characterization file.

	Example:
		from anadama2 import Workflow
		from fugassem.tools.generate_annotation_input import prepare_annotation

		# create an anadama2 workflow instance
		workflow=Workflow()

		# add prepare_function tasks
		specific_ann_file = prepare_co_unit (
						ann_file,
						ann_type,
						unit_file,
						func_file,
						output_folder,
						extracted_ann_file,
						workflow,
                        args.threads,
                        args.time_equation,
                        args.mem_equation,
						func_type)
		# run the workflow
		workflow.go()
	"""

	config.logger.info("###### Start prepare_co_unit module #####")

	# prep I/O files
	main_folder = output_folder
	if not os.path.isdir(main_folder):
		os.system("mkdir -p " + main_folder)
	extracted_ann_file = os.path.join(main_folder, os.path.basename(extracted_ann_file))
	extracted_log = extracted_ann_file + ".log"

	workflow.add_task(
		"fugassem_prepare_seqSimilarity -a [depends[0]] -s [depends[1]] -f [depends[2]] -t [args[0]] -m [args[1]] -o [targets[0]] > [args[2]] 2>&1",
		depends = [ann_file, unit_file, func_file, TrackedExecutable("fugassem_prepare_seqSimilarity")],
		targets = [extracted_ann_file],
		args = [ann_type, func_type, extracted_log],
		cores = 1,
		time = time_equation,
		mem = mem_equation,
		name = "fugassem_prepare_seqSimilarity")

	return extracted_ann_file


def prepare_annotation_task (ann_file, func_type, pfam, ddi,
                             contig, clust_file, gene_info_file,
							 homology, homology_ann_file, coanns,
                             output_folder, basename,
                             workflow, threads, time_equation, mem_equation):
	"""
	Prepare annotation info used for FUGAsseM inputs

	Args:
		ann_file: MetaWIBELE finalized characterization file.
		func_type: function type, [Default: GO ]
		pfam: if specified, will extract pfam annotation of that type, [ Default: "InterProScan_PfamDomain" ]
		ddi: if specified, will extract DDI annotation of that type, [ Default: "DOMINE_interaction" ]
		contig: if specified, will extract contig annotation, [ Default: None ]
		clust_file: clustering file in extended-fasta format for protein families, which can be generated by MetaWIBELE [ Default: None ]
		gene_info_file: gene info file including the contig info of each member in protein families [ Default: None ]
		homology: if specified, will extract co-homology annotation, [ Default: None ]
		homology_ann_file: homology-based annotation file [ Default: None ]
		coanns: a list of data types which are converted into co-annotation pairs
		output_folder (string): The path of the output folder.
		basename: basename for output files.
		workflow (anadama2.workflow): An instance of the workflow class.
		threads (int): The number of threads/cores for clustering to use.
		time_equation (int): requred number of hours defined in the workflow.
		mem_equation (int): requred number of GB defined in the workflow.

	Requires:
		annotation files

	Returns:
		string: annotations inputs used for FUGAsseM

	Example:
		from anadama2 import Workflow
		from fugassem.tools.generate_annotation_input import prepare_annotation_info

		# create an anadama2 workflow instance
		workflow=Workflow()

		# add prepare_function tasks
		 prepare_annotation_task (
						ann_file,
						func_type,
						pfam,
						ddi,
                        contig,
                        clust_file,
                        gene_info_file,
						homology,
						homology_ann_file,
						coanns,
                        output_folder,
                        basename
						workflow,
                        args.threads,
                        args.time_equation,
                        args.mem_equation)
		# run the workflow
		workflow.go()
	"""

	config.logger.info("###### Start prepare_annotation_task module #####")

	# prep I/O files
	main_folder = output_folder
	if not os.path.isdir(main_folder):
		os.system("mkdir -p " + main_folder)

	# function annotations
	funcs = func_type.split(",")
	ann_files = []
	func_files = []
	func_types = []
	for mytype in funcs:
		myfunc1 = os.path.join(main_folder, basename + "." + mytype + ".tsv")
		myfunc2 = os.path.join(main_folder, basename + "." + mytype + ".simple.tsv")
		prepare_annotation(ann_file, mytype, None,
		                   main_folder, myfunc1, myfunc2,
		                   workflow, threads, time_equation, mem_equation)
		func_files.append(myfunc2)
		func_types.append(mytype)

	# pfam
	if pfam:
		myann1 = os.path.join(main_folder, basename + ".pfam.tsv")
		myann2 = os.path.join(main_folder, basename + ".pfam.simple.tsv")
		prepare_annotation(ann_file, pfam, None,
		                   main_folder, myann1, myann2,
		                   workflow, threads, time_equation, mem_equation)
		if "pfam" in coanns:
			myann2_raw = myann2
			myann2 = os.path.join(main_folder, basename + ".pfam.simple.pairs.tsv")
			prepare_coann_pairs (myann2_raw,
			                    main_folder, myann2,
			                    workflow, threads, time_equation, mem_equation)
		ann_files.append(myann2)

	# ddi
	if ddi:
		myann1 = os.path.join(main_folder, basename + ".DDI.tsv")
		myann2 = os.path.join(main_folder, basename + ".DDI.simple.tsv")
		prepare_annotation(ann_file, ddi, None,
		                   main_folder, myann1, myann2,
		                   workflow, threads, time_equation, mem_equation)
		if "ddi" in coanns:
			myann2_raw = myann2
			myann2 = os.path.join(main_folder, basename + ".DDI.simple.pairs.tsv")
			prepare_coann_pairs (myann2_raw,
			                    main_folder, myann2,
			                    workflow, threads, time_equation, mem_equation)
		ann_files.append(myann2)

	# contig
	if contig:
		myann1 = os.path.join(main_folder, basename + ".contig.tsv")
		myann2 = os.path.join(main_folder, basename + ".contig.simple.tsv")
		prepare_contig(clust_file, gene_info_file,
		               output_folder, myann1, myann2,
		               workflow, threads, time_equation, mem_equation)
		if "contig" in coanns:
			myann2_raw = myann2
			myann2 = os.path.join(main_folder, basename + ".contig.simple.pairs.tsv")
			prepare_coann_pairs (myann2_raw,
			                    main_folder, myann2,
			                    workflow, threads, time_equation, mem_equation)
		ann_files.append(myann2)

	# homology
	if homology:
		for myfunc in func_files:
			mytype = func_types[func_files.index(myfunc)]
			mym = os.path.basename(myfunc)
			mygo = mym.split(".")[-3]
			myann = os.path.join(main_folder, basename + "." + mygo + ".homology.tsv")
			prepare_co_unit(ann_file, "homology",
		                homology_ann_file, myfunc,
		                output_folder, myann,
		                workflow, threads, time_equation, mem_equation, mytype)
			ann_files.append(myann)

	return ann_files, func_files


def generate_annotation_input (workflow):
	'''
	Prepare annotations used as inputs for FUGAsseM
	'''

	## get arguments
	args = workflow.parse_args()

	## prep settings
	if args.threads:
		try:
			args.threads = int(args.threads)
		except:
			config.logger.info ("Please provide valid number of required threads! Otherwise, will set it as 1")
			args.threads = 1
	else:
		args.threads = 1
	if args.memory:
		try:
			args.memory = int(args.memory)
		except:
			config.logger.info("Please provide valid number of required memory! Otherwise, will set it as 10240")
			args.memory = 10240
	else:
		args.memory = 10240
	if args.time:
		try:
			args.time = int(args.time)
		except:
			config.logger.info("Please provide valid number of required time! Otherwise, will set it as 600")
			args.time = 600
	else:
		args.time = 600
	if args.coann_list:
		coanns = args.coann_list.split(",")
	else:
		coanns = []

	## get all input files
	ann_file = os.path.abspath(args.input)
	if args.basename:
		basename = args.basename
	else:
		basename = "fugassem"
	if not os.path.isfile(ann_file):
		sys.exit("Please provide annotation file!")

	## prepare outputs
	output_dir = os.getcwd()
	if args.output:
		output_dir = os.path.abspath(args.output)
	if not os.path.isdir(output_dir):
		os.system("mkdir -p " + output_dir)

	## run prepare-annotation_info module
	config.logger.info("Start to run prepare_annotation_task module......")
	prepare_annotation_task (ann_file, args.func_type, args.pfam, args.ddi,
	                        args.contig, args.clust_file, args.gene_info,
	                        args.homology, args.homology_ann, coanns,
	                        output_dir, basename,
	                        workflow, args.threads, args.time, args.memory)

	## start the workflow
	workflow.go()


def main():
	workflow = parse_cli_arguments()
	generate_annotation_input (workflow)


if __name__ == "__main__":
	main()
