#!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 import filter_HGT
from MetaCHIP.MetaCHIP_config import config_dict


to_do = '''

1. if sh: FastTree: command not found exit, as well as for all other programs
2. combine grouping and group_to_taxon file
3. not print if disabled: [2018-12-09 21:54:32] Plotting flanking regions with 16 cores
4. steps move to PG.py: uclust, get species tree
5. make a tmp folder
6. 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 :::...
        
    Core modules:
       PI             ->    Prepare input files 
       BP             ->    Run Best-match and Phylogenetic approaches
       
    Supplementary modules:
       CMLP           ->    Combine multiple level predictions (part of BP module)
       filter_HGT     ->    Get HGTs predicted at least n levels (for multiple level predictions)
       sankey_taxon   ->    Show taxonomic classification with Sankey plot [to be added]
       get_SCG_tree   ->    Get SCG protein tree [to be added]
       plot_circos    ->    Plot circos [to be added]
       
    # for command specific help
    MetaCHIP PI -h
    MetaCHIP BP -h

    ''' % version()

    print(help_message)


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: 90')
    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')

    # arguments for CMLP
    CMLP_parser = subparsers.add_parser('CMLP', description='Combine multiple level predictions', epilog='Example: MetaCHIP CMLP -h')
    CMLP_parser.add_argument('-p',                      required=True,                          help='output prefix')
    CMLP_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')
    CMLP_parser.add_argument('-cov',                    required=False, type=int,   default=75, help='coverage cutoff, default: 75')
    CMLP_parser.add_argument('-al',                     required=False, type=int,   default=200,help='alignment length cutoff, default: 200')
    CMLP_parser.add_argument('-flk',                    required=False, type=int,   default=10, help='the length of flanking sequences to plot (Kbp), default: 10')
    CMLP_parser.add_argument('-ip',                     required=False, type=int,   default=90, help='identity percentile cutoff, default: 90')
    CMLP_parser.add_argument('-ei',                     required=False, type=float, default=90, help='end match identity cutoff, default: 90')
    CMLP_parser.add_argument('-t',                      required=False, type=int,   default=1,  help='number of threads, default: 1')

    # arguments for filter_HGT
    CMLP_parser = subparsers.add_parser('filter_HGT', description='Combine multiple level predictions', epilog='Example: MetaCHIP CMLP -h')
    CMLP_parser.add_argument('-i',                      required=True,                          help='txt file containing detected HGTs, e.g. [prefix]_[ranks]_detected_HGTs.txt ')
    CMLP_parser.add_argument('-n',                      required=True, type=int,                help='HGTs detected at least n levels, between 1 - 5')
    CMLP_parser.add_argument('-plot',                   required=False,                         help='flanking plots folder')

    # 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())

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


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

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

        detection_ranks_str = args['r']

        # 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':

        detection_ranks_str = args['r']

        # 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)

    #################### run supplementary modules ####################

    if args['subparser_name'] == 'CMLP':
        BP.combine_multiple_level_predictions(args, config_dict)

    if args['subparser_name'] == 'filter_HGT':
        filter_HGT.filter_HGT(args)
