#!/usr/bin/env python3

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

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

# MarkerMAG 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 datetime import datetime
from MarkerMAG.MarkerMAG_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 = ''' 
         ...::: MarkerMAG v%s :::...

    link             ->  link MAGs with marker genes
    rename_reads     ->  rename paired reads
    matam_16s        ->  reconstruct 16S from metagenome with Matam
    uclust_16s       ->  cluster 16S genes with Usearch
    barrnap_16s      ->  identify 16S genes with Barrnap
    subsample_reads  ->  subsample reads with Usearch

    ''' % version(config_dict)

    print(help_message)
    #     get_16s_depth    ->  get 16S depth
    #     get_mag_depth    ->  get MAG depth


if __name__ == '__main__':

    # 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] == 'link':
        from MarkerMAG import link_16s
        link_16s_parser = subparsers.add_parser('link', description='Linking MAGs with 16S rRNA marker genes', usage=link_16s.link_Marker_MAG_usage)
        default_prefix = 'MyRun_%s' % datetime.now().strftime('%Y-%m-%d_%H_%M_%S')

        #############################
        # specify argument group
        link_16s_parser_input_files = link_16s_parser.add_argument_group("input files")
        link_16s_parser_16s         = link_16s_parser.add_argument_group("16S rRNA gene related parameters")
        link_16s_parser_both_rds    = link_16s_parser.add_argument_group("parameters for linking (round 1 and 2)")
        link_16s_parser_rd1         = link_16s_parser.add_argument_group("parameters for linking (round 1)")
        link_16s_parser_rd2         = link_16s_parser.add_argument_group("parameters for linking (round 2)")
        link_16s_parser_preset      = link_16s_parser.add_argument_group("preset parameters, decide automatically if not specified")
        link_16s_parser_dependency  = link_16s_parser.add_argument_group("provide if dependencies are not in your system path")
        link_16s_parser_others      = link_16s_parser.add_argument_group("program settings")
        link_16s_parser_debug       = link_16s_parser.add_argument_group("for debugging, do NOT specify")

        # input files
        link_16s_parser_input_files.add_argument('-p',          required=False, metavar='',             default=default_prefix, help='output prefix, (default: MyRun_SystemTime)')
        link_16s_parser_input_files.add_argument('-r1',         required=True,  metavar='',                                     help='paired reads r1 (fasta format)')
        link_16s_parser_input_files.add_argument('-r2',         required=True,  metavar='',                                     help='paired reads r2 (fasta format)')
        link_16s_parser_input_files.add_argument('-marker',     required=True,  metavar='',                                     help='marker gene sequences')
        link_16s_parser_input_files.add_argument('-mag',        required=True,  metavar='',                                     help='metagenome-assembled-genome (MAG) folder')
        link_16s_parser_input_files.add_argument('-x',          required=False, metavar='',             default='fasta',        help='MAG file extension, (default: %(default)s)')

        # 16S rRNA gene related parameters
        link_16s_parser_16s.add_argument('-min_len_16s',        required=False, metavar='', type=int,   default=1200,           help='minimum length of 16S (in bp) for linking, (default: %(default)s)')
        link_16s_parser_16s.add_argument('-no_polish',          required=False, action="store_true",                            help='skip polishing 16S before linking')
        link_16s_parser_16s.add_argument('-min_iden_16s',       required=False, metavar='', type=float, default=98.5,           help='minimum similarity for 16S sequences to be assigned to the same genome, (default: %(default)s)')
        link_16s_parser_16s.add_argument('-min_cov_16s',        required=False, metavar='', type=float, default=50,             help='coverage cutoff for calculating pairwise 16S similarity, (default: %(default)s)')
        link_16s_parser_16s.add_argument('-min_aln_16s',        required=False, metavar='', type=int,   default=500,            help='alignment length cutoff for calculating pairwise 16S similarity, (default: %(default)s)')
        link_16s_parser_16s.add_argument('-keep_ctg_end_16s',   required=False, action="store_true",                            help='do not mask 16S sequences at the end of MAG contig')

        # parameters for both rounds linking
        link_16s_parser_both_rds.add_argument('-mismatch',      required=False, metavar='', type=float, default=2,              help='maximum mismatch percentage, (default: %(default)s)')
        link_16s_parser_both_rds.add_argument('-min_M_len_16s', required=False, metavar='', type=int,   default=45,             help='minimum length aligned to 16S, (default: %(default)s)')
        link_16s_parser_both_rds.add_argument('-min_M_len_ctg', required=False, metavar='', type=int,   default=45,             help='minimum length aligned to ctg, (default: %(default)s)')
        link_16s_parser_both_rds.add_argument('-min_M_pct',     required=False, metavar='', type=float, default=35,             help='minimum aligned percentage, (default: %(default)s)')
        link_16s_parser_both_rds.add_argument('-min_link',      required=False, metavar='', type=int,   default=8,              help='minimum number of linkages to report, (default: %(default)s)')
        link_16s_parser_both_rds.add_argument('-link_num_diff', required=False, metavar='', type=float, default=80,             help='within_gnm_linkage_num_diff, (default: %(default)s)')

        # parameters for 2nd round linking
        link_16s_parser_rd2.add_argument('-mira',               required=False, action="store_true",                            help='run Mira, instead of Spades')
        link_16s_parser_rd2.add_argument('-mira_tmp',           required=False, default=None,                                   help='tmp dir for mira')

        # program settings
        link_16s_parser_others.add_argument('-t',               required=False, metavar='', type=int,   default=1,              help='number of threads, (default: %(default)s)')
        link_16s_parser_others.add_argument('-tmp',             required=False, action="store_true",                            help='keep temporary files')
        link_16s_parser_others.add_argument('-quiet',           required=False, action="store_true",                            help='not report progress')
        link_16s_parser_others.add_argument('-force',           required=False, action="store_true",                            help='force overwrite existing results')

        # for debugging
        link_16s_parser_debug.add_argument('-test_mode',        required=False, action="store_true",                            help='only for debugging, do not provide')
        link_16s_parser_debug.add_argument('-rd2_only',         required=False, action="store_true",                            help='run round 2 only')
        link_16s_parser_debug.add_argument('-sorted_sam16s',    required=False, default=None,                                   help='mapping of input reads to 16S')
        link_16s_parser_debug.add_argument('-vis_all',          required=False, action="store_true",                            help='vis all linkages, including those filtered out, for debugging')

        args = vars(parser.parse_args())
        link_16s.link_16s(args)

    elif sys.argv[1] == 'rename_reads':
        from MarkerMAG import rename_reads
        rename_reads_parser = subparsers.add_parser('rename_reads', description='Rename reads', usage=rename_reads.rename_reads_usage)
        rename_reads_parser.add_argument('-r1',  required=True, type=str,               help='forward reads, fasta format')
        rename_reads_parser.add_argument('-r2',  required=True, type=str,               help='reverse reads, fasta format')
        rename_reads_parser.add_argument('-p',   required=True, type=str,               help='prefix of output file and read id')
        rename_reads_parser.add_argument('-fq',  required=False, action="store_true",   help='rename reads with quality score')
        rename_reads_parser.add_argument('-t',   required=False, type=int, default=1,   help='number of threads, default: 1')
        args = vars(parser.parse_args())
        rename_reads.rename_reads(args)

    elif sys.argv[1] == 'matam_16s':
        from MarkerMAG import matam_16s
        matam_16s_parser = subparsers.add_parser('matam_16s', description='Assemble 16S rRNA genes with Matam', usage=matam_16s.matam_16s_usage)
        matam_16s_parser.add_argument('-p',              required=True,  metavar='',                                            help='output prefix')
        matam_16s_parser.add_argument('-r1',             required=False, metavar='', default=None,                              help='forward reads')
        matam_16s_parser.add_argument('-r2',             required=False, metavar='', default=None,                              help='reverse reads')
        matam_16s_parser.add_argument('-r16s',           required=False, metavar='', default=None,                              help='extracted 16S reads, specify if you already have 16S reads identified from metagenomic dataset')
        matam_16s_parser.add_argument('-pct',            required=True,  metavar='', type=str, default='1,5,10,25,50,75,100',   help='subsample percentage, must be integer, between 1-100, deafault: 1,5,10,25,50,75,100')
        matam_16s_parser.add_argument('-d',              required=False, metavar='', type=str,                                  help='Matam ref db, controls "-d" in Matam')
        matam_16s_parser.add_argument('-i',              required=False, metavar='', type=float, default=0.999,                 help='cluster identity cutoff, do not provide value lower than 0.99, default: 0.999')
        matam_16s_parser.add_argument('-t',              required=False, metavar='', type=int, default=1,                       help='number of threads, default: 1')
        matam_16s_parser.add_argument('-mem',            required=False, metavar='', type=int, default=10240,                   help='maximum memory to be allocated to Matam (in MB), controls "--max_memory" in Matam, default: 10240')
        matam_16s_parser.add_argument('-force',          required=False, action="store_true",                                   help='force overwrite existing results')
        matam_16s_parser.add_argument('-quiet',          required=False, action="store_true",                                   help='not report progress')
        matam_16s_parser.add_argument('-matam',          required=False, metavar='', type=str, default='matam_assembly.py',     help='path to matam_assembly.py, default: matam_assembly.py')
        matam_16s_parser.add_argument('-seqtk',          required=False, metavar='', type=str, default='seqtk',                 help='path to seqtk executable file, default: seqtk')
        matam_16s_parser.add_argument('-usearch',        required=False, metavar='', type=str, default='usearch',               help='path to usearch executable file, default: usearch')
        args = vars(parser.parse_args())
        matam_16s.matam_16s(args)

    elif sys.argv[1] == 'uclust_16s':
        from MarkerMAG import uclust_16s
        uclust_16S_parser = subparsers.add_parser('uclust_16s', description='Cluster 16S with Usearch', usage=uclust_16s.uclust_16S_usage)
        uclust_16S_parser.add_argument('-in',           required=True,                          help='fasta file')
        uclust_16S_parser.add_argument('-i',            required=False, type=float, default=1,  help='Identity cutoff (0-1), default: 1')
        uclust_16S_parser.add_argument('-out',          required=True,                          help='file out')
        args = vars(parser.parse_args())
        uclust_16s.uclust_16S(args)

    elif sys.argv[1] == 'subsample_reads':
        from MarkerMAG import subsample_reads
        subsample_reads_parser = subparsers.add_parser('subsample_reads', description='subsample paired reads', usage=subsample_reads.subsample_reads_usage)
        subsample_reads_parser.add_argument('-r1',      required=True, type=str,            help='forward reads')
        subsample_reads_parser.add_argument('-r2',      required=True, type=str,            help='reverse reads')
        subsample_reads_parser.add_argument('-ratio',   required=True, type=str,            help='subsample ratio, 0-1')
        subsample_reads_parser.add_argument('-usearch', required=False, default='usearch',  help='path to usearch executable file, default: usearch')
        args = vars(parser.parse_args())
        subsample_reads.subsample_reads(args)

    elif sys.argv[1] == 'barrnap_16s':
        from MarkerMAG import barrnap_16s
        barrnap_16s_parser = subparsers.add_parser('barrnap_16s', description='identify 16S sequences with Barrnap', usage=barrnap_16s.barrnap_16s_usage)
        barrnap_16s_parser.add_argument('-p',       required=True, type=str,               help='output prefix')
        barrnap_16s_parser.add_argument('-g',       required=True, type=str,               help='genome folder')
        barrnap_16s_parser.add_argument('-x',       required=False, default='fasta',       help='genome file extension, default: fasta')
        barrnap_16s_parser.add_argument('-t',       required=False, type=int, default=1,   help='number of threads, default: 1')
        barrnap_16s_parser.add_argument('-force',   required=False, action="store_true",   help='force overwrite existing results')
        args = vars(parser.parse_args())
        barrnap_16s.barrnap_16s(args)

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


upload_to_pypi_cmd = '''

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

shan88

#twine upload --repository-url https://test.pypi.org/legacy/ dist/*

pip3 install --upgrade MarkerMAG
pip3 install --upgrade -i https://test.pypi.org/simple/ MarkerMAG

'''
