#!python

# Copyright (C) 2026, Shan Zhang and Weizhi Song (songwz03@gmail.com).
# MetaCHIP2 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.

# MetaCHIP2 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
import warnings
from MetaCHIP2.MetaCHIP2_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 = ''' 
           ...::: MetaCHIP2 v%s :::...
    
    Core module
       detect   ->  Detect HGT
    
    Supplementary modules
       tree     ->  Infer genome tree
       prokka   ->  A wrapper for running Prokka in batch manner
       circos   ->  Visualize HGTs with circos plot
       stats    ->  (to be added) Get stats of predicted HGTs
       enrich   ->  Enrichment of COG functions in HGTs
           
    # Upgrade with: pip3 install --upgrade MetaCHIP2
    ''' % 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')
    warnings.filterwarnings('ignore')

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

    elif sys.argv[1] == 'detect':
        from MetaCHIP2 import detect
        detect_parser = subparsers.add_parser('detect', description='detect HGT', usage=detect.detect_usage)
        detect_parser.add_argument('-o',        required=True,                         help='output folder')
        detect_parser.add_argument('-i',        required=True,                         help='input gbk folder')
        detect_parser.add_argument('-x',        required=False, default='gbk',         help='file extension, default: gbk')
        detect_parser.add_argument('-s',        required=True,                         help='species tree of input genomes')
        detect_parser.add_argument('-c',        required=False, default=None,          help='taxonomic classification of input genomes')
        detect_parser.add_argument('-r',        required=False, default=None,          help='grouping rank, e.g., p, c, o, f, g, pcofg or pco...')
        detect_parser.add_argument('-g',        required=False, default=None,          help='grouping file')
        detect_parser.add_argument('-b',        required=False, default=None,          help='all-vs-all blastn results (e.g., from a previous run)')
        detect_parser.add_argument('-m',        required=False, action="store_true",   help='speed-up all-vs-all blastn using mmseqs')
        detect_parser.add_argument('-np',       required=False, action="store_true",   help='skip plotting flanking regions')
        detect_parser.add_argument('-nc',       required=False, action="store_true",   help='skip end-break and contig-match check, not recommend for metagenome-assembled genomes')
        detect_parser.add_argument('-t',        required=False, type=int, default=1,   help='number of threads, default: 1')
        detect_parser.add_argument('-tmp',      required=False, action="store_true",   help='keep temporary files')
        detect_parser.add_argument('-q',        required=False, action="store_true",   help='do not report progress')
        detect_parser.add_argument('-f',        required=False, action="store_true",   help='force overwrite previous results')
        args = vars(parser.parse_args())
        detect.detect(args)

    elif sys.argv[1] == 'prokka':
        from MetaCHIP2 import prokka
        prokka_parser = subparsers.add_parser('prokka', description='run prokka', usage=prokka.prokka_usage)
        prokka_parser.add_argument('-i',          required=True,                          help='genome folder')
        prokka_parser.add_argument('-x',          required=True, default='fna',           help='file extension, deafult: fna')
        prokka_parser.add_argument('-d',          required=False, default=None,           help='genome domain, Bacteria or Archaea')
        prokka_parser.add_argument('-m',          required=False, action="store_true",    help='annotate MAG')
        prokka_parser.add_argument('-e',          required=False, action="store_true",    help='execute commands')
        prokka_parser.add_argument('-t',          required=False, type=int, default=1,    help='number of threads')
        args = vars(parser.parse_args())
        prokka.prokka(args)

    elif sys.argv[1] == 'circos':
        from MetaCHIP2 import circos
        circos_parser = subparsers.add_parser('circos', description='circos', usage=circos.circos_usage)
        circos_parser.add_argument('-i',   required=False, default=None,         help='input matrix')
        circos_parser.add_argument('-l',   required=False, default=None,         help='MetaCHIP produced detected_HGTs.txt')
        circos_parser.add_argument('-g',   required=False, default=None,         help='grouping file')
        circos_parser.add_argument('-o',   required=True,                        help='output plot')
        args = vars(parser.parse_args())
        circos.circos(args)

    elif sys.argv[1] == 'tree':
        from MetaCHIP2 import tree
        tree_parser = subparsers.add_parser('tree', description='Infer species tree of input genomes', usage=tree.tree_usage)
        tree_parser.add_argument('-i',  required=True,                          help='genome folder')
        tree_parser.add_argument('-x',  required=False, default='fna',          help='genome file extension, default: fna')
        tree_parser.add_argument('-c',  required=False, default=None,           help='taxonomic classification of input genomes')
        tree_parser.add_argument('-o',  required=True,                          help='output folder')
        tree_parser.add_argument('-db', required=True,                          help='GTDB database files')
        tree_parser.add_argument('-t',  required=False, type=int, default=1,    help='number of threads')
        tree_parser.add_argument('-f',  required=False, action="store_true",    help='force overwrite existing results')
        args = vars(parser.parse_args())
        tree.tree(args)

    elif sys.argv[1] == 'enrich':
        from MetaCHIP2 import enrich
        enrich_parser = subparsers.add_parser('enrich', description='enrich', usage=enrich.enrich_usage)
        enrich_parser.add_argument('-o',       required=True,                              help='output plot')
        enrich_parser.add_argument('-faa',     required=True,                              help='folder of faa files, produced by Prokka')
        enrich_parser.add_argument('-x',       required=False, default='faa',              help='file extension, default: faa')
        enrich_parser.add_argument('-hgt1',    required=True,                              help='amino acid sequences of HGTs, required')
        enrich_parser.add_argument('-hgt2',    required=False, default=None,               help='amino acid sequences of HGTs, e.g., predicted with a different approach, optional')
        enrich_parser.add_argument('-db',      required=True,                              help='COG_db_dir')
        enrich_parser.add_argument('-diamond', required=False, action='store_true',        help='run diamond (for big dataset), default is NCBI blastp')
        enrich_parser.add_argument('-t',       required=False, type=int, default=1,        help='number of threads')
        enrich_parser.add_argument('-e',       required=False, default=0.001, type=float,  help='evalue cutoff, default: 0.001')
        enrich_parser.add_argument('-desc',    required=False, action='store_true',        help='include functional description in x-axis labels')
        enrich_parser.add_argument('-f',       required=False, action="store_true",        help='force overwrite')
        args = vars(parser.parse_args())
        enrich.enrich(args)

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


upload_to_pypi_cmd = '''

cd /Users/songweizhi/PycharmProjects/MetaCHIP2
rm -r build dist MetaCHIP2.egg-info
python3 setup.py sdist bdist_wheel
twine upload dist/* --verbose

__token__

pip3 install --upgrade MetaCHIP2        

'''

'''
When encounter the following error, try to replace name="MetaCHIP2" to name="metachip2" in the setup.py. 
HTTPError: 400 Bad Request from https://upload.pypi.org/legacy/
Bad Request  
'''