#!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 argparse
from MetaCHIP.PI import PI
from MetaCHIP.BM import BM
from MetaCHIP.PG import PG
#from MetaCHIP.PG_MultiRanks import PG_MultiRanks
from MetaCHIP import MetaCHIP_config
from MetaCHIP.MetaCHIP_config import config_dict


to_do = '''

All:
!!!!!! predict HGT at levels from phylum to genus
add option -force
replace all-vs-all blastn with usearch
if sh: FastTree: command not found exit
add hgt_wf: PI, BM and PG


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


BM:
not print if disabled: [2018-12-09 21:54:32] Plotting flanking regions with 16 cores


PG:
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 
       BM               ->   Best-Match approach 
       PG               ->   PhyloGenetic approach

    # for command specific help
    MetaCHIP <command> -h
    ''' % version()

    print(help_message)

'''
    Tree manipulation:
       plot_tree        ->   [to be added] Plot tree in newick format
       subset_gtdb_tree ->   [to be added] Subset GTDB tree for provided taxons
       
    Annotation modules (to be added):
       run_prodigal     ->   [to be added] Run Prodigal
       cog_annotation   ->   [to be added] Perform COG annotation
       run_barrnap      ->   [to be added] Run barrnap
       get_SCG_tree     ->   [to be added] Get SCG protein tree
       
    Plot modules (to be added):
       plot_circos      ->   [to be added] Plot circos
       plot_identity    ->   [to be added] Plot identity distribution
       get_gene_cluster ->   [to be added] Get gene clusters with Usearch
       plot_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')
    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('-qsub',          required=False, action="store_true", help='run blastn with job scripts, only for HPC users')
    PI_parser.add_argument('-force',         required=False, action="store_true", help='overwrite previous results')
    PI_parser.add_argument('-quiet',         required=False, action="store_true", help='not report progress')

    # arguments for BM approach
    BM_parser = subparsers.add_parser('BM',  description='Best-match approach', epilog='Example: MetaCHIP BM -h')
    BM_parser.add_argument('-p',             required=True,  help='output prefix')
    BM_parser.add_argument('-r',             required=False, default=None, help='grouping rank')
    BM_parser.add_argument('-g',             required=False, default=None, help='grouping file')
    BM_parser.add_argument('-cov',           required=False, type=int, default=75, help='coverage cutoff, default: 75')
    BM_parser.add_argument('-al',            required=False, type=int, default=200, help='alignment length cutoff, default: 200')
    BM_parser.add_argument('-flk',           required=False, type=int, default=10, help='the length of flanking sequences to plot (Kbp), default: 10')
    BM_parser.add_argument('-ip',            required=False, type=int, default=90, help='identity percentile cutoff, default: 90')
    BM_parser.add_argument('-ei',            required=False, type=float, default=90, help='end match identity cutoff, default: 95')
    BM_parser.add_argument('-t',             required=False, type=int, default=1, help='number of threads, default: 1')
    BM_parser.add_argument('-plot_iden',     required=False, action="store_true", help='plot identity distribution')
    BM_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)')
    BM_parser.add_argument('-force',         required=False, action="store_true", help='overwrite previous results')
    BM_parser.add_argument('-quiet',         required=False, action="store_true", help='Do not report progress')
    BM_parser.add_argument('-tmp',           required=False, action="store_true", help='keep temporary files')

    # arguments for PG approach
    PG_parser = subparsers.add_parser('PG',  description='Phylogenetic approach', epilog='Example: MetaCHIP PG -h')
    PG_parser.add_argument('-p',             required=True,  help='output prefix')
    PG_parser.add_argument('-r',             required=False, default=None, help='grouping rank')
    PG_parser.add_argument('-g',             required=False, help='grouping file')
    PG_parser.add_argument('-cov',           required=False, type=int, default=75, help='coverage cutoff, default: 75')
    PG_parser.add_argument('-al',            required=False, type=int, default=200, help='alignment length cutoff, default: 200')
    PG_parser.add_argument('-flk',           required=False, type=int, default=10, help='the length of flanking sequences to plot (Kbp), default: 10')
    PG_parser.add_argument('-ip',            required=False, type=int, default=90, help='identity percentile, default: 90')
    PG_parser.add_argument('-ei',            required=False, type=float, default=90, help='end match identity cutoff, default: 95')
    PG_parser.add_argument('-t',             required=False, type=int, default=1, help='number of threads, default: 1')
    PG_parser.add_argument('-force',         required=False, action="store_true", help='overwrite previous results')
    PG_parser.add_argument('-quiet',         required=False, action="store_true", help='Do not report progress')


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

    # run corresponding module
    if args['subparser_name'] == 'PI':
        PI(args, config_dict)

    if args['subparser_name'] == 'BM':
        BM(args, config_dict)

    if args['subparser_name'] == 'PG':
        PG(args, config_dict)
