#!/usr/bin/env python

# Copyright (C) 2017, Weizhi Song, Torsten Thomas.
# songwz03@gmail.com or t.thomas@unsw.edu.au

# MetaCHIP is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# MetaCHIP is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import sys
import copy
import argparse
from MetaCHIP.PI import PI
from MetaCHIP import BP
from datetime import datetime
from MetaCHIP import MetaCHIP_config
from MetaCHIP.MetaCHIP_config import config_dict


to_do = '''

All:
add option -force
replace all-vs-all blastn with usearch
if sh: FastTree: command not found exit, as well as for all other programs
add hgt_wf: PI, BM and PG
combine NorthSea_c5_grouping.txt and NorthSea_c5_group_to_taxon.txt into one file

PI:
if there is a Usearch error, break the pipeline
combine grouping and group_to_taxon file

BP:
not print if disabled: [2018-12-09 21:54:32] Plotting flanking regions with 16 cores
change text direction for circos plot
steps move to PG.py: uclust, get species tree
removed %s columns from the concatenated msa with low coverage 
removed %s columns from the concatenated msa with consensus
make a tmp folder
if no PG validated, skip plot, but not report error

'''


def version():
    version_file = open('%s/VERSION' % MetaCHIP_config.config_file_path)
    return version_file.readline().strip()


def print_main_help():

    help_message = ''' 
        ...::: MetaCHIP v%s :::...
        
    HGT detection modules:
       PI      ->      Prepare Input files 
       BP      ->      Best-match and Phylogenetic approach

    # for command specific help
    MetaCHIP PI -h
    MetaCHIP BP -h

    ''' % version()

    print(help_message)


'''

To be added:       
    get_SCG_tree     ->   [to be added] Get SCG protein tree
    plot_circos      ->   [to be added] Plot circos
    sankey_taxon       ->   [to be added] Show taxon correlations with sankey plot
'''


if __name__ == '__main__':

    # initialize the options parser
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(help="--", dest='subparser_name')

    # arguments for PI
    PI_parser = subparsers.add_parser('PI',     description='Prepare input files', epilog='Example: MetaCHIP PI -h')
    PI_parser.add_argument('-i',                required=True,                       help='input genome folder')
    PI_parser.add_argument('-taxon',            required=False,                      help='taxonomic classification')
    PI_parser.add_argument('-p',                required=True,                       help='output prefix')
    PI_parser.add_argument('-r',                required=False, default=None,        help='grouping rank, choose from p (phylum), c (class), o (order), f (family) or g (genus)')
    PI_parser.add_argument('-g',                required=False, default=None,        help='grouping file')
    PI_parser.add_argument('-x',                required=False, default='fasta',     help='file extension')
    PI_parser.add_argument('-grouping_only',    required=False, action="store_true", help='run grouping only, deactivate Prodigal and Blastn')
    PI_parser.add_argument('-nonmeta',          required=False, action="store_true", help='annotate Non-metagenome-assembled genomes (Non-MAGs)')
    PI_parser.add_argument('-noblast',          required=False, action="store_true", help='not run all-vs-all blastn')
    PI_parser.add_argument('-t',                required=False, type=int, default=1, help='number of threads, default: 1')
    PI_parser.add_argument('-blastn_js_header', required=False,                      help='speed up all-against-all blastn with separated job script for each of the input genome, provide the job script header here')
    PI_parser.add_argument('-qsub',             required=False, action="store_true", help='specify to automatically submit generated job scripts, otherwise, submit them manually')
    PI_parser.add_argument('-quiet',            required=False, action="store_true", help='not report progress')

    # arguments for BP approach
    BP_parser = subparsers.add_parser('BP',     description='BM and PG approach', epilog='Example: MetaCHIP BP -h')
    BP_parser.add_argument('-p',                required=True,                          help='output prefix')
    BP_parser.add_argument('-r',                required=False, default=None,           help='grouping rank, choose from p (phylum), c (class), o (order), f (family), g (genus) or any combination of them')
    BP_parser.add_argument('-g',                required=False, default=None,           help='grouping file')
    BP_parser.add_argument('-cov',              required=False, type=int,   default=75, help='coverage cutoff, default: 75')
    BP_parser.add_argument('-al',               required=False, type=int,   default=200,help='alignment length cutoff, default: 200')
    BP_parser.add_argument('-flk',              required=False, type=int,   default=10, help='the length of flanking sequences to plot (Kbp), default: 10')
    BP_parser.add_argument('-ip',               required=False, type=int,   default=90, help='identity percentile cutoff, default: 90')
    BP_parser.add_argument('-ei',               required=False, type=float, default=90, help='end match identity cutoff, default: 95')
    BP_parser.add_argument('-t',                required=False, type=int,   default=1,  help='number of threads, default: 1')
    BP_parser.add_argument('-plot_iden',        required=False, action="store_true",    help='plot identity distribution')
    BP_parser.add_argument('-NoEbCheck',        required=False, action="store_true",    help='disable end break and contig match check for fast processing, not recommend for metagenome-assembled genomes (MAGs)')
    BP_parser.add_argument('-force',            required=False, action="store_true",    help='overwrite previous results')
    BP_parser.add_argument('-quiet',            required=False, action="store_true",    help='Do not report progress')
    BP_parser.add_argument('-tmp',              required=False, action="store_true",    help='keep temporary files')


    # get and check options
    args = None
    if (len(sys.argv) == 1) or (sys.argv[1] == '-h') or (sys.argv[1] == '-help') or (sys.argv[1] == '--help'):
        print_main_help()
        sys.exit(0)

    else:
        args = vars(parser.parse_args())

    detection_ranks_str = args['r']
    time_format = '[%Y-%m-%d %H:%M:%S]'


    #################### run PI module ####################

    if args['subparser_name'] == 'PI':

        # for single level detection
        if len(detection_ranks_str) == 1:
            PI(args, config_dict)

        # for multiple level prediction
        if len(detection_ranks_str) > 1:
            gene_predicted = 0
            for detection_rank_PI in detection_ranks_str:
                current_rank_args_PI = copy.deepcopy(args)
                current_rank_args_PI['r'] = detection_rank_PI

                if gene_predicted == 0:
                    current_rank_args_PI['grouping_only'] = False
                    PI(current_rank_args_PI, config_dict)
                    gene_predicted = 1
                else:
                    current_rank_args_PI['grouping_only'] = True
                    PI(current_rank_args_PI, config_dict)


    #################### run BP module ####################

    if args['subparser_name'] == 'BP':

        # for single level detection
        if len(detection_ranks_str) == 1:
            BP.BM(args, config_dict)
            BP.PG(args, config_dict)

        # for multiple level prediction
        if len(detection_ranks_str) > 1:
            for detection_rank_BP in detection_ranks_str:
                current_rank_args_BP = copy.deepcopy(args)
                current_rank_args_BP['r'] = detection_rank_BP
                current_rank_args_BP['quiet'] = True

                print('%s Detect HGT at level: %s' % ((datetime.now().strftime(time_format)), detection_rank_BP))
                BP.BM(current_rank_args_BP, config_dict)
                BP.PG(current_rank_args_BP, config_dict)

        # combine multiple level predictions
        BP.combine_multiple_level_predictions(args, config_dict)

