"""
FUGAsseM utility: prepare stratified-MTX file as FUGAsseM input
A collection of tasks for preparing stratified-MTX file that will be used as FUGAsseM input

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 FUGAsseM
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 stratified MTX abundance",
	                    remove_options=["output"])

	# add the custom arguments to the workflow
	workflow.add_argument("extension-paired",
	                      desc = "provide the extension for paired fastq files using comma to separate, e.g. .R1.fastq.gz,.R2.fastq.gz | .R1.fastq,.R2.fastq",
	                      default = ".R1.fastq.gz,.R2.fastq.gz")
	workflow.add_argument("extension",
	                      desc = "provide the extension for all fastq files",
	                      choices = [".fastq.gz", ".fastq"],
	                      default = ".fastq.gz")
	workflow.add_argument("taxon-level",
	                      desc = "taxonomic level used for stratification [ Default: Species ]",
	                      choices = ["MSP", "Species", "Terminal"],
	                      default = "Species")
	workflow.add_argument("taxon-prevalence",
	                      desc = "minimum prevalence of each taxon [ Default: None ]",
	                      default = None)
	workflow.add_argument("normalized-method",
	                      desc = "normalization scheme: copies per million [cpm], relative abundance [relab]; [ Default: cpm ]",
	                      default = "cpm")
	workflow.add_argument("gene-catalog",
	                      desc = "input clustering file in extended-fasta format for non-redundant gene catalogs, which can be generated by MetaWIBELE",
	                      required = True)
	workflow.add_argument("gene-catalog-seq",
	                      desc = "input fasta file of nucleotide sequences of non-redundant gene catalogs, which can be generated by MetaWIBELE",
	                      required = True)
	workflow.add_argument("protein-family",
	                      desc = "input clustering file in extended-fasta format for protein families clustered by non-redundant gene catalogs, which can be generated by MetaWIBELE",
	                      required = True)
	workflow.add_argument("family-taxonomy",
	                      desc = "input taxonomy file for protein families, which can be generated by MetaWIBELE",
	                      required = True)
	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 get_gene_abundance(input_dir, extension, extension_paired, gene_catalog, gene_catalog_nuc,
                       output_folder, gene_catalog_count,
                       workflow, threads, time_equation, mem_equation):
	"""
    This set of tasks used for calculating MTX abundance.

    Args:
		input_dir: path of FastQ files
		extension: provide the extension for all fastq files [default: .fastq.gz]
		extension_paired: provide the extension for paired fastq files using comma to separate, e.g. ".R1.fastq.gz,.R2.fastq.gz" [ default: None ]
		gene_catalog: cluster file formatted in extended-fasta of non-redundant gene catalogs.
		gene_catalog_nuc: fastq file of nucleotide sequences for non-redundant gene catalogs.
		output_folder: folder name for outputs.
		gene_catalog_count: count file for gene non-redundant gene catalogs.
		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:
        bowtie2 (tested with 2.3.2)
        samtools (tested with 1.5)
        featureCounts (tested with Version 1.6.2)
        the nucleotide sequences of non-redundant gene catalogs
        fastq files for each sample

    Returns:
        string: file of count data for non-redundant gene catalogs

    Example:
        from anadama2 import Workflow
		from fugassem.tools.generate_stratified_mtx_input import get_gene_abundance

        # create an anadama2 workflow instance
        workflow=Workflow()

        # add quality control tasks for the fastq files
		mycounts = get_gene_abundance (input_dir, extension, extension_paired, gene_catalog, gene_catalog_nuc,
						output_folder, gene_catalog_count,
                        workflow, threads, time_equation, mem_equation)

        # run the workflow
        workflow.go()
    """

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

	## prep I/O files
	main_folder = output_folder
	if not os.path.isdir(main_folder):
		os.system("mkdir -p " + main_folder)
	gene_catalog_count = os.path.join(main_folder, os.path.basename(gene_catalog_count))
	gene_catalog_saf = re.sub(".fna$", ".saf.gtf", gene_catalog_nuc)
	gene_catalog_saf = re.sub(".fasta$", ".saf.gtf", gene_catalog_saf)
	#gene_catalog_saf = gene_catalog_nuc + ".saf.gtf"
	base_name = os.path.basename(gene_catalog_nuc)
	if re.search("\.([^\.]+)$", base_name):
		tmp = re.search("\.([^\.]+)$", base_name)
		tmp = tmp.group(1)
		prefix_gene_catalog = re.sub("." + tmp + "$", "", base_name)
	else:
		prefix_gene_catalog = base_name

	mylog = gene_catalog_saf + ".log"
	workflow.add_task(
		'fugassem_gene_abundance_indexRef -r [depends[0]] -t gene -b [args[0]] -o [targets[0]] >[args[1]] 2>&1 ',
		depends=[gene_catalog_nuc, TrackedExecutable("fugassem_gene_abundance_indexRef")],
		targets=[gene_catalog_saf],
		args=[prefix_gene_catalog, mylog],
		cores=1,
		name="fugassem_gene_abundance_indexRef")

	## collect sequences
	if extension_paired:
		extension_paireds = extension_paired.split(",")
		sample_files = utilities.find_files(input_dir, extension_paireds[0], None)
		samples = utilities.sample_names(sample_files, extension_paireds[0], None)
	else:
		sample_files = utilities.find_files(input_dir, extension, None)
		samples = utilities.sample_names(sample_files, extension, None)

	## bowtie2 will map reads to gene categories
	flt_seqs = []
	for sample in samples:
		seq_file = "NA"
		if extension_paired:
			tmp = extension_paired.split(",")
		else:
			if extension != "none":
				tmp = extension.split(",")
		for item in tmp:
			if seq_file == "NA":
				seq_file = os.path.join(input_dir, sample + '%s' % item)
			else:
				seq_file = seq_file + "," + os.path.join(input_dir, sample + '%s' % item)
		flt_seqs.append((sample, seq_file))
	# foreah sample

	## Now run bowtie2 to map reads to gene categories
	mappings = []
	mappings_tmp = []
	for (sample, seq_file) in flt_seqs:
		seq_base = sample
		mydir = os.path.join(output_folder, sample)
		os.system("mkdir -p " + mydir)
		sample_counts = os.path.join(mydir, seq_base + ".sort.bed")
		stdout_log = os.path.join(mydir, '%s.mapping.stdout.log' % seq_base)
		mappings_tmp.append(sample_counts)

		workflow.add_task(
			'fugassem_gene_abundance -r [depends[0]] -u [args[0]] -t [args[1]] -s [args[2]] -w [args[3]] '
			'> [args[4]] 2>&1 ',
			depends=[gene_catalog_nuc, gene_catalog_saf, TrackedExecutable("fugassem_gene_abundance")],
			targets=[sample_counts],
			args=[seq_file, threads, seq_base, mydir, stdout_log],
			cores=1,
			name=sample + "__gene_abundance")

	for myfile in mappings_tmp:
		myname = os.path.basename(myfile)
		myfile_new = os.path.join(output_folder, myname)
		mappings.append(myfile_new)
		workflow.add_task(
			"ln -fs [depends[0]] [targets[0]]",
			depends=[myfile],
			targets=[myfile_new],
			cores=1,
			name="ln__" + myname)

	# collect abundance
	mylog = gene_catalog_count + ".log"
	workflow.add_task(
		'fugassem_gene_catalog_abundance -p [args[0]] -s sort.bed -c [args[1]] -o [targets[0]] >[args[2]] 2>&1 ',
		depends=utilities.add_to_list(mappings, TrackedExecutable("fugassem_gene_catalog_abundance")),
		targets=[gene_catalog_count],
		args=[output_folder, gene_catalog, mylog],
		cores=1,
		name="fugassem_gene_catalog_abundance")

	return gene_catalog_count


def stratify_mtx_abundance (input_dir, extension, extension_paired, gene_catalog_file, gene_catalog_nuc,
                           protein_family_file, protein_family_tax_file,
                           taxon_level, taxon_prevalence, normalized_method,
                           output_folder, gene_catalog_count, stratified_abund_file,
                           workflow, threads, time_equation, mem_equation):
	"""
	Generate stratified MTX abundance used as input for FUGAsseM

	Args:
		input_dir: path of FastQ files
		extension: provide the extension for all fastq files [default: .fastq.gz]
		extension_paired: provide the extension for paired fastq files using comma to separate, e.g. ".R1.fastq.gz,.R2.fastq.gz" [ default: None ]
		gene_catalog: cluster file formatted in extended-fasta of non-redundant gene catalogs.
		gene_catalog_nuc: fastq file of nucleotide sequences for non-redundant gene catalogs.
		proterin_clust_file: cluster file formatted in extended-fasta of protein families that are clustered from non-redundant gene catalogs.
		proterin_family_tax_file: taxonomy file for protein families that are clustered from non-redundant gene catalogs.
		taxon_level: stratified taxnomic level, [choices: "Species", "MPS"] [ Default: Species]
		taxon_prevalence: threshold for taxon prevalence (e.g. >X% of genes non-zero values) [ Default: None]
		normalized_method: normalization scheme: copies per million [cpm], relative abundance [relab]; default=[cpm]
		output_folder: folder name for outputs.
		gene_catalog_count: count file for gene non-redundant gene catalogs.
		stratified_abund_file: stratified MTX abundance file.
		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:
		fastq files for each sample
		the nucleotide sequences of non-redundant gene catalogs
        the protein families clustered by non-redundant gene catalogs

	Returns:
		string: the file of stratified MTX abundance file

	Example:
		from anadama2 import Workflow
		from fugassem.tools.generate_stratified_mtx_input import stratify_mtx_abundance

		# create an anadama2 workflow instance
		workflow=Workflow()

		# add preprocess_function tasks
		stratified_mtx_file = stratify_mtx_abundance (
							input_dir, extension, extension_paired, gene_catalog_file, gene_catalog_nuc,
                            protein_family_file, protein_family_tax_file,
                            taxon_level, taxon_prevalence, normalized_method,
							output_folder, gene_catalog_count, stratified_abund_file,
                            workflow, threads, time_equation, mem_equation)
		# run the workflow
		workflow.go()
	"""

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

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

	# get gene counts
	get_gene_abundance(input_dir, extension, extension_paired, gene_catalog_file, gene_catalog_nuc,
	                   output_folder, gene_catalog_count,
	                   workflow, threads, time_equation, mem_equation)

	# collect count for protein families
	family_count_all = re.sub(".tsv", ".counts.tsv", stratified_abund_file)
	mylog = family_count_all + ".log"
	workflow.add_task_gridable(
		"fugassem_sum_to_family_abundance -i [depends[0]] -c [depends[1]] -o [targets[0]] >[args[0]] 2>&1",
		depends=[gene_catalog_count, protein_family_file, TrackedExecutable("fugassem_sum_to_family_abundance")],
		targets=[family_count_all],
		args=[mylog],
		cores=1,
		time=time_equation,
		mem=mem_equation,
		name="fugassem_sum_to_family_abundance")

	# normalize to RPK
	family_rpk = re.sub(".tsv", ".RPK.tsv", stratified_abund_file)
	mylog = re.sub(".tsv", ".log", family_rpk)
	workflow.add_task(
		"fugassem_abundance_RPK -i [depends[0]] -l [depends[1]] -t cluster -o [targets[0]] >[args[0]] 2>&1",
		depends=[family_count_all, protein_family_file, TrackedExecutable("fugassem_abundance_RPK")],
		targets=[family_rpk],
		args=[mylog],
		cores=1,
		name="fugassem_abundance_RPK")

	# stratify abundance
	family_stratify = re.sub(".tsv", ".stratified_RPK.tsv", stratified_abund_file)
	mylog = family_stratify + ".log"
	workflow.add_task(
		"fugassem_stratified_abundance -i [depends[0]] -m [depends[1]] -r [args[0]] -t [args[1]] -o [targets[0]] >[args[2]] 2>&1",
		depends=[family_rpk, protein_family_tax_file, TrackedExecutable("fugassem_stratified_abundance")],
		targets=[family_stratify],
		args=[taxon_level, taxon_prevalence, mylog],
		cores=1,
		name="fugassem_stratified_abundance")

	# normalize within taxon
	mylog = stratified_abund_file + ".log"
	workflow.add_task(
		"fugassem_abundance_normalization -i [depends[0]] -u [args[0]] -m [args[1]] -o [targets[0]] >[args[2]] 2>&1",
		depends=[family_stratify, TrackedExecutable("fugassem_abundance_normalization")],
		targets=[stratified_abund_file],
		args=[normalized_method, "taxonwise", mylog],
		cores=1,
		name="fugassem_abundance_normalization")

	return stratified_abund_file


def generate_stratified_mtx_input (workflow):
	'''
	Prepare stratified MTX abundance used as input 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

	## get all input files
	input_dir = os.path.abspath(args.input)
	if args.gene_catalog:
		gene_catalog = os.path.abspath(args.gene_catalog)
	if args.gene_catalog_seq:
		gene_catalog_seq = os.path.abspath(args.gene_catalog_seq)
	if args.protein_family:
		protein_family = os.path.abspath(args.protein_family)
	if args.family_taxonomy:
		family_taxonomy = os.path.abspath(args.family_taxonomy)
	if args.basename:
		basename = args.basename
	if not os.path.isfile(gene_catalog):
		sys.exit("Please provide gene-catalog file!")
	if not os.path.isfile(gene_catalog_seq):
		sys.exit("Please provide gene-catalog nucleotide sequence file!")
	if not os.path.isfile(protein_family):
		sys.exit("Please provide protein-family file!")
	if not os.path.isfile(family_taxonomy):
		sys.exit("Please provide taxonomy file for protein families!")

	## 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)
	gene_catalog_seq_raw = gene_catalog_seq
	gene_catalog_seq = os.path.join(output_dir, os.path.basename(gene_catalog_seq_raw))
	os.system("cp " + gene_catalog_seq_raw + " " + gene_catalog_seq)
	gene_catalog_count = os.path.join(output_dir, basename + ".genecatalogs.counts.tsv")
	stratified_abund_file = os.path.join(output_dir, basename + ".proteinfamilies.nrm.tsv")

	## run stratified-mtx-abundance module
	config.logger.info("Start to run stratified-mtx-abundance module......")
	stratify_mtx_abundance (input_dir, args.extension, args.extension_paired, gene_catalog, gene_catalog_seq,
	                           protein_family, family_taxonomy,
	                           args.taxon_level, args.taxon_prevalence, args.normalized_method,
	                           output_dir, gene_catalog_count, stratified_abund_file,
	                           workflow, args.threads, args.time, args.memory)

	## start the workflow
	workflow.go()


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


if __name__ == "__main__":
	main()
