#!/usr/bin/env python3

# Copyright (C) 2017, Weizhi Song.
# songwz03@gmail.com

# TreeSAK 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.

# TreeSAK 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 warnings
import argparse
from TreeSAK.TreeSAK_config import config_dict


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


def print_main_help():

    help_message = ''' 
                 ...::: TreeSAK v%s :::...
                 
    Marker-related
       parse_deltall_stdout   ->  Parse stdout of deltaLL.rb
       get_arCOG_seq          ->  Retrieve arCOG sequences
       AssessMarkerPA         ->  Assess Markers by P/A among groups
    
    Multiple Sequence Alignment
       ConvertMSA             ->  Convert MSA format
       fa2phy                 ->  Convert MSA format (fasta to phylip)
       OneLineAln             ->  One-line fasta format alignments
       SliceMSA               ->  Slice MSA by column 
    
    Tree-related
       GTDB_tree_r207         ->  Infer GTDB (r207) archaeal/bacterial tree
       subset_tree            ->  Subset tree
       compare_trees          ->  Compare trees with Mantel test
       rename_leaves          ->  Rename tree leaves
       FLN                    ->  Format leaf names (e.g. remove spaces in names)
    
    Dating-related
       AssessCVG              ->  Assess dating convergence
    
    Dating workflow
       Marker2Tree            ->  Marker to Tree
       AssessMarkerDeltaLL    ->  Assess Markers by DeltaLL
       Dating                 ->  Perform dating
    
    ''' % version(config_dict)

    print(help_message)


if __name__ == '__main__':

    ########################################################################################### initialize subparsers ############################################################################################

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

    # disable warning message
    warnings.filterwarnings('ignore')

    # parse options
    if (len(sys.argv) == 1) or (sys.argv[1] in ['-h', '-help', '--help']):
        print_main_help()
        sys.exit(0)

    elif sys.argv[1] == 'subset_tree':
        from TreeSAK import subset_tree
        subset_tree_parser = subparsers.add_parser('subset_tree', description='Subset tree', usage=subset_tree.subset_tree_usage)
        subset_tree_parser.add_argument('-tree',            required=True,                                  help='input tree file')
        subset_tree_parser.add_argument('-taxon',           required=True,                                  help='A file containing list of nodes to keep, one node per line')
        subset_tree_parser.add_argument('-out',             required=True,                                  help='Output tree file')
        subset_tree_parser.add_argument('-q',               required=False, action="store_true",            help='do not report progress')
        args = vars(parser.parse_args())
        subset_tree.subset_tree(args)

    elif sys.argv[1] == 'label_tree':
        from TreeSAK import label_tree
        label_tree_parser = subparsers.add_parser('label_tree', description='Add labels to tree leaves', usage=label_tree.label_tree_usage)
        label_tree_parser.add_argument('-tree',             required=True,                                  help='tree file in newick format')
        label_tree_parser.add_argument('-label',            required=False,  default=None,                  help='label file (label,leaf)')
        label_tree_parser.add_argument('-taxon',            required=False,  default=None,                  help='taxonomic classification')
        label_tree_parser.add_argument('-rank',             required=False,  default=None,                  help='taxonomic rank to label')
        args = vars(parser.parse_args())
        label_tree.label_tree(args, config_dict)

    elif sys.argv[1] == 'OneLineAln':
        from TreeSAK import OneLineAln
        OneLineAln_parser = subparsers.add_parser('OneLineAln', description='One-line fasta format alignments', usage=OneLineAln.OneLineAln_usage)
        OneLineAln_parser.add_argument('-in',       required=True,                       help='input MSA in fasta format')
        OneLineAln_parser.add_argument('-out',      required=False, default=None,        help='output file')
        OneLineAln_parser.add_argument('-upper',    required=False, action='store_true', help='turn to uppercase')
        args = vars(parser.parse_args())
        OneLineAln.OneLineAln(args)

    elif sys.argv[1] == 'SliceMSA':
        from TreeSAK import SliceMSA
        SliceMSA_parser = subparsers.add_parser('SliceMSA', description='Slice MSA by column', usage=SliceMSA.SliceMSA_usage)
        SliceMSA_parser.add_argument('-i',   required=True,                         help='input MSA in fasta format')
        SliceMSA_parser.add_argument('-s',   required=True,                         help='columns to export, e.g. 200-300, -100, 50-')
        SliceMSA_parser.add_argument('-o',   required=True,                         help='output folder')
        SliceMSA_parser.add_argument('-f',   required=False, action="store_true",   help='force overwrite existing output folder')
        args = vars(parser.parse_args())
        SliceMSA.SliceMSA(args)

    elif sys.argv[1] == 'compare_trees':
        from TreeSAK import compare_trees
        compare_trees_parser = subparsers.add_parser('compare_trees', usage=compare_trees.compare_trees_usage)
        compare_trees_parser.add_argument('-t1',  required=True,                       help='tree (folder) 1')
        compare_trees_parser.add_argument('-t2',  required=True,                       help='tree (folder) 2')
        compare_trees_parser.add_argument('-tx',  required=False, default='newick',    help='extention of tree files, default: newick')
        compare_trees_parser.add_argument('-dm', required=False, action="store_true",  help='export distance-alike matrix, obtained by subtract the similarity value from 1')
        compare_trees_parser.add_argument('-t',   required=False, type=int, default=1, help='number of threads')
        compare_trees_parser.add_argument('-tmp', required=False, action="store_true", help='keep tmp files')
        args = vars(parser.parse_args())
        compare_trees.compare_trees(args)

    elif sys.argv[1] == 'rename_leaves':
        from TreeSAK import rename_leaves
        rename_leaves_parser = subparsers.add_parser('rename_leaves', usage=rename_leaves.rename_leaves_usage)
        rename_leaves_parser.add_argument('-i',    required=True,             help='input tree')
        rename_leaves_parser.add_argument('-r',    required=True,             help='rename file')
        rename_leaves_parser.add_argument('-f',    required=False, default=1, help='tree format, default: 1')
        rename_leaves_parser.add_argument('-o',    required=True,             help='output tree')
        args = vars(parser.parse_args())
        rename_leaves.rename_leaves(args)

    elif sys.argv[1] == 'FLN':
        from TreeSAK import format_leaf_name
        format_leaf_name_parser = subparsers.add_parser('FLN', usage=format_leaf_name.format_leaf_name_usage)
        format_leaf_name_parser.add_argument('-i',                  required=True,                          help='input tree')
        format_leaf_name_parser.add_argument('-fmt',                required=False, default=1,              help='tree format, default: 1')
        format_leaf_name_parser.add_argument('-o',                  required=True,                          help='output tree')
        format_leaf_name_parser.add_argument('-s2u',                required=False, action="store_true",    help='change space in tree leaves to underscore')
        format_leaf_name_parser.add_argument('-ns',                 required=False, action="store_true",    help='remove space from leaf names')
        format_leaf_name_parser.add_argument('-nsqm',               required=False, action="store_true",    help='remove single quotation marks from leaf names')
        format_leaf_name_parser.add_argument('-ndqm',               required=False, action="store_true",    help='remove double quotation marks from leaf names')
        args = vars(parser.parse_args())
        format_leaf_name.format_leaf_name(args)

    elif sys.argv[1] == 'AssessCVG':
        from TreeSAK import AssessCVG
        AssessCVG_parser = subparsers.add_parser('AssessCVG', usage=AssessCVG.AssessCVG_usage)
        AssessCVG_parser.add_argument('-m1', required=True, help='mcmc.txt from run 1')
        AssessCVG_parser.add_argument('-m2', required=True, help='mcmc.txt from run 2')
        AssessCVG_parser.add_argument('-o',  required=True, help='output convergence plot')
        args = vars(parser.parse_args())
        AssessCVG.AssessCVG(args)

    elif sys.argv[1] == 'parse_deltall_stdout':
        from TreeSAK import parse_deltall_stdout
        parse_deltall_stdout_parser = subparsers.add_parser('parse_deltall_stdout', usage=parse_deltall_stdout.parse_deltall_stdout_usage)
        parse_deltall_stdout_parser.add_argument('-i', required=True, help='input file (e.g., nohup.out)')
        parse_deltall_stdout_parser.add_argument('-o', required=True, help='output summary')
        args = vars(parser.parse_args())
        parse_deltall_stdout.parse_deltall_stdout(args)

    elif sys.argv[1] == 'get_arCOG_seq':
        from TreeSAK import get_arCOG_seq
        get_arCOG_seq_parser = subparsers.add_parser('get_arCOG_seq', usage=get_arCOG_seq.get_arCOG_seq_usage)
        get_arCOG_seq_parser.add_argument('-i',      required=True,                         help='arCOD id file, one id per line')
        get_arCOG_seq_parser.add_argument('-db_dir', required=True,                         help='database folder')
        get_arCOG_seq_parser.add_argument('-o',      required=True,                         help='output folder')
        get_arCOG_seq_parser.add_argument('-f',      required=False, action="store_true",   help='force overwrite existing output folder')
        args = vars(parser.parse_args())
        get_arCOG_seq.get_arCOG_seq(args)

    elif sys.argv[1] == 'ConvertMSA':
        from TreeSAK import ConvertMSA
        ConvertMSA_parser = subparsers.add_parser('ConvertMSA', usage=ConvertMSA.ConvertMSA_usage)
        ConvertMSA_parser.add_argument('-i',       required=True,                       help='input alignment')
        ConvertMSA_parser.add_argument('-xi',      required=False, default='aln',       help='input alignment extension')
        ConvertMSA_parser.add_argument('-fi',      required=True,                       help='input alignment format, e.g., fasta, phylip')
        ConvertMSA_parser.add_argument('-o',       required=True,                       help='output alignment')
        ConvertMSA_parser.add_argument('-xo',      required=False, default='aln',       help='output alignment extension')
        ConvertMSA_parser.add_argument('-fo',      required=True,                       help='output alignment format, e.g., fasta, phylip')
        ConvertMSA_parser.add_argument('-oneline', required=False, action="store_true", help='put sequence in single line, available if -fo is fasta')
        ConvertMSA_parser.add_argument('-nogap',   required=False, action="store_true", help='remove gaps from alignment, available if -fo is fasta')
        ConvertMSA_parser.add_argument('-f',       required=False, action="store_true", help='force overwrite existing output folder')
        args = vars(parser.parse_args())
        ConvertMSA.ConvertMSA(args)

    elif sys.argv[1] == 'Marker2Tree':
        from TreeSAK import Marker2Tree
        Marker2Tree_parser = subparsers.add_parser('Marker2Tree', usage=Marker2Tree.Marker2Tree_usage)
        Marker2Tree_parser.add_argument('-m',               required=True,                          help='marker seq dir')
        Marker2Tree_parser.add_argument('-mx',              required=True,                          help='marker seq ext')
        Marker2Tree_parser.add_argument('-aa',              required=True,                          help='faa file dir')
        Marker2Tree_parser.add_argument('-aax',             required=True,                          help='faa file ext')
        Marker2Tree_parser.add_argument('-g',               required=True,                          help='genome group')
        Marker2Tree_parser.add_argument('-pac',             required=False, default='0-50-75-100',  help='cutoffs, default: 0-50-75-100')
        Marker2Tree_parser.add_argument('-o',               required=True,                          help='output dir')
        Marker2Tree_parser.add_argument('-e',               required=True,                          help='e-value')
        Marker2Tree_parser.add_argument('-t',               required=True,  type=int,               help='num of threads')
        Marker2Tree_parser.add_argument('-skip_align_trim', required=False, action="store_true",    help='skip extracting, aligning and trimming markers')
        Marker2Tree_parser.add_argument('-mmn',             required=False, default=10, type=int,   help='minimal marker number, default: 10')
        Marker2Tree_parser.add_argument('-jst',             required=False, default='6',            help='threads to request in job script, for running iqtree')
        Marker2Tree_parser.add_argument('-qsub',            required=False, action="store_true",    help='submit job scripts')
        Marker2Tree_parser.add_argument('-f',               required=False, action="store_true",    help='force overwrite')
        args = vars(parser.parse_args())
        Marker2Tree.Marker2Tree(args)

    elif sys.argv[1] == 'AssessMarkerPA':
        from TreeSAK import AssessMarkerPA
        AssessMarkerPA_parser = subparsers.add_parser('AssessMarkerPA', usage=AssessMarkerPA.AssessMarkerPA_usage)
        AssessMarkerPA_parser.add_argument('-ta',         required=True,                               help='trimmed alignments')
        AssessMarkerPA_parser.add_argument('-tax',        required=True,                               help='extension of trimmed alignments')
        AssessMarkerPA_parser.add_argument('-aa',         required=True,                               help='faa file dir')
        AssessMarkerPA_parser.add_argument('-aax',        required=True,                               help='faa file ext')
        AssessMarkerPA_parser.add_argument('-g',          required=True,                               help='genome group')
        AssessMarkerPA_parser.add_argument('-c',          required=False, default='25-50-75-85-100',   help='cutoffs, default: 25-50-75-85-100')
        AssessMarkerPA_parser.add_argument('-o',          required=True,                               help='output dir')
        AssessMarkerPA_parser.add_argument('-f',          required=False, action="store_true",         help='force overwrite existing output folder')
        args = vars(parser.parse_args())
        AssessMarkerPA.AssessMarkerPA(args)

    elif sys.argv[1] == 'AssessMarkerDeltaLL':
        from TreeSAK import AssessMarkerDeltaLL
        AssessMarkerDeltaLL_parser = subparsers.add_parser('AssessMarkerDeltaLL', usage=AssessMarkerDeltaLL.AssessMarkerDeltaLL_usage)
        AssessMarkerDeltaLL_parser.add_argument('-deltall',     required=True,                          help='DeltaLL stdout')
        AssessMarkerDeltaLL_parser.add_argument('-o',           required=True,                          help='output dir')
        AssessMarkerDeltaLL_parser.add_argument('-c',           required=False, default='25-50-75-100', help='cutoffs, default: 25-50-75-100')
        AssessMarkerDeltaLL_parser.add_argument('-mmn',         required=False, default=20, type=int,   help='minimal marker number, default: 20')
        AssessMarkerDeltaLL_parser.add_argument('-aln',         required=True,                          help='faa file dir')
        AssessMarkerDeltaLL_parser.add_argument('-jst',         required=False, default='6',            help='threads to request in job script, for running iqtree')
        AssessMarkerDeltaLL_parser.add_argument('-qsub',        required=False, action="store_true",    help='submit job scripts')
        AssessMarkerDeltaLL_parser.add_argument('-f',           required=False, action="store_true",    help='force overwrite')
        args = vars(parser.parse_args())
        AssessMarkerDeltaLL.AssessMarkerDeltaLL(args)

    elif sys.argv[1] == 'Dating':
        from TreeSAK import Dating
        Dating_parser = subparsers.add_parser('Dating', usage=Dating.Dating_usage)
        Dating_parser.add_argument('-deltall', required=True,                          help='DeltaLL stdout')
        Dating_parser.add_argument('-aod',     required=True,                          help='AssessMarkerDeltaLL output dir')
        Dating_parser.add_argument('-og',      required=True,                          help='outgroup leaves, one leaf id per line')
        Dating_parser.add_argument('-eu',      required=True,                          help='EU tree with time constraints')
        Dating_parser.add_argument('-o',       required=True,                          help='dating wd')
        Dating_parser.add_argument('-c',       required=False, default='25-50-75-100', help='cutoffs, default: 25-50-75-100')
        Dating_parser.add_argument('-mmn',     required=False, default=20, type=int,   help='minimal marker number, default: 20')
        Dating_parser.add_argument('-ra',      required=False, default=45, type=int,   help='root age, default: 45')
        Dating_parser.add_argument('-jst',     required=False, default='6',            help='threads to request in job script, for performing dating')
        Dating_parser.add_argument('-qsub',    required=False, action="store_true",    help='submit job scripts for getting in.BV')
        Dating_parser.add_argument('-f',       required=False, action="store_true",    help='force overwrite')
        args = vars(parser.parse_args())
        Dating.Dating(args)

    elif sys.argv[1] == 'fa2phy':
        from TreeSAK import fa2phy
        fa2phy_parser = subparsers.add_parser('fa2phy', usage=fa2phy.fa2phy_usage)
        fa2phy_parser.add_argument('-i', required=True, help='input MSA in fasta format')
        fa2phy_parser.add_argument('-o', required=True, help='output MSA in phylip format')
        args = vars(parser.parse_args())
        fa2phy.fa2phy(args)

    else:
        print('Unrecognized command: %s, program exited' % sys.argv[1])
        exit()


upload_to_pypi_cmd = '''
cd /Users/songweizhi/PycharmProjects/TreeSAK
rm -r build dist TreeSAK.egg-info
python setup.py sdist bdist_wheel
twine upload dist/*

songweizhi  shan88
pip3 install --upgrade TreeSAK
'''

