#!/usr/bin/env python3

###############################################################################
#                                                                             #
#    This program is free software: you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation, either version 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program 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 General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program. If not, see <http://www.gnu.org/licenses/>.     #
#                                                                             #
###############################################################################

__author__ = "Donovan Parks"
__copyright__ = "Copyright 2017"
__credits__ = ["Donovan Parks"]
__license__ = "GPL3"
__maintainer__ = "Donovan Parks"
__email__ = "donovan.parks@gmail.com"
__status__ = "Development"

import sys
import argparse

from unitem import version
from unitem.main import OptionsParser
from unitem.utils import CustomHelpFormatter, logger_setup


def print_help():
    """Help menu."""

    print('')
    print('                ...::: UniteM v' + version() + ' :::...''')
    print('''\

  Binning:
    bin -> Apply binning methods to an assembly
  
  Profiling:
    profile -> Identify marker genes and calculate assembly statistics
   
  Bin selection:
    greedy    -> Greedy bin selection across multiple binning methods
    consensus -> Consensus clustering across multiple binning methods
    unanimous -> Unanimous bin filtering across multiple binning methods

  Use: unitem <command> -h for command specific help.
  Version: unitem -v

  Feature requests or bug reports can be sent to Donovan Parks (donovan.parks@gmail.com)
    or posted on GitHub (https://github.com/dparks1134/unitem).
    ''')


if __name__ == '__main__':

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

    # bin
    bin_parser = subparsers.add_parser('bin',
                                       formatter_class=CustomHelpFormatter,
                                       description='Apply binning methods to an assembly.',
                                       add_help=False)

    bin_parser.add_argument(
        'assembly_file', help="assembled contigs to bin (FASTA format)")
    bin_parser.add_argument('output_dir', help="output directory")

    required_group = bin_parser.add_argument_group(
        title='required arguments (one of)')
    cov_group = required_group.add_mutually_exclusive_group(required=True)
    cov_group.add_argument('--bam_files', nargs='*',
                           help="BAM file(s) to parse for coverage profile")
    cov_group.add_argument('--cov_file',
                           help="file indicating coverage information")

    binning_group = bin_parser.add_argument_group(title='binning methods')
    binning_group.add_argument('--gm2',
                               help="run GroopM v2",
                               action='store_true')
    binning_group.add_argument('--max40',
                               help="run MaxBin with the 40 marker gene set",
                               action='store_true')
    binning_group.add_argument('--max107',
                               help="run MaxBin with the 107 marker gene set",
                               action='store_true')
    binning_group.add_argument('--mb2',
                               help="run MetaBAT v2",
                               action='store_true')
    binning_group.add_argument('--mb_verysensitive',
                               help="run MetaBAT v1 with the 'verysensitive' preset settings",
                               action='store_true')
    binning_group.add_argument('--mb_sensitive',
                               help="run MetaBAT v1 with the 'sensitive' preset settings",
                               action='store_true')
    binning_group.add_argument('--mb_specific',
                               help="run MetaBAT v1 with the 'specific' preset settings",
                               action='store_true')
    binning_group.add_argument('--mb_veryspecific',
                               help="run MetaBAT v1 with the 'veryspecific' preset settings",
                               action='store_true')
    binning_group.add_argument('--mb_superspecific',
                               help="run MetaBAT v1 with the 'superspecific' preset settings",
                               action='store_true')

    optional_group = bin_parser.add_argument_group(title='optional arguments')
    optional_group.add_argument('-m', '--min_contig_len',
                                help="minimum length of contigs to bin (>=1500 recommended)",
                                type=int, default=2500)
    optional_group.add_argument('-c', '--cpus',
                                help="number of CPUs",
                                type=int, default=1)
    optional_group.add_argument('--silent',
                                help="suppress output of logger",
                                action='store_true')
    optional_group.add_argument('-h', '--help',
                                help="show help message",
                                action="help")

    # profile bins
    profile_parser = subparsers.add_parser('profile',
                                           formatter_class=CustomHelpFormatter,
                                           description='Identify marker genes and calculate assembly statistics.',
                                           add_help=False)

    profile_parser.add_argument('output_dir', help="output directory")

    required_group = profile_parser.add_argument_group(
        title='required arguments (one of)')
    bin_group = required_group.add_mutually_exclusive_group(required=True)
    bin_group.add_argument('-b', '--bin_dirs',
                           help="directories with bins from different binning methods", nargs='*')
    bin_group.add_argument('-f', '--bin_file',
                           help="tab-separated file indicating directories with bins from binning methods (two columns: method name and directory)")

    optional_group = profile_parser.add_argument_group(
        title='optional arguments')
    optional_group.add_argument('--marker_dir',
                                help="directory containing Pfam and TIGRfam marker genes data")
    optional_group.add_argument('-c', '--cpus',
                                help="number of CPUs",
                                type=int, default=1)
    optional_group.add_argument('--silent',
                                help="suppress output of logger",
                                action='store_true')
    optional_group.add_argument('-h', '--help',
                                help="show help message",
                                action="help")

    # consensus bin selection
    consensus_parser = subparsers.add_parser('consensus',
                                             formatter_class=CustomHelpFormatter,
                                             description="Consensus clustering across multiple binning methods.",
                                             add_help=False)

    consensus_parser.add_argument(
        'profile_dir', help="directory with bin profiles (output of 'profile' command)")
    consensus_parser.add_argument('output_dir', help="output directory")

    required_group = consensus_parser.add_argument_group(
        title='required arguments (one of)')
    bin_group = required_group.add_mutually_exclusive_group(required=True)
    bin_group.add_argument('-b', '--bin_dirs',
                           help="directories with bins from different binning methods", nargs='*')
    bin_group.add_argument('-f', '--bin_file',
                           help="tab-separated file indicating directories with bins from binning methods (two columns: method name and directory)")

    consensus_group = consensus_parser.add_argument_group(
        title="optional bin selection arguments")
    consensus_group.add_argument('-w', '--weight', help="weight given to contamination for assessing genome quality",
                                 type=float, default=2)
    consensus_group.add_argument('-q', '--sel_min_quality', help="minimum quality of bin to consider during bin selection process",
                                 type=float, default=50)
    consensus_group.add_argument('-x', '--sel_min_comp', help="minimum completeness of bin to consider during bin selection process",
                                 type=float, default=50)
    consensus_group.add_argument('-y', '--sel_max_cont', help="maximum contamination of bin to consider during bin selection process",
                                 type=float, default=10)
    consensus_group.add_argument('-r', '--remove_perc', help="minimum percentage of bins required to remove contigs from highest-quality bin",
                                 type=float, default=50.0)
    consensus_group.add_argument('-a', '--add_perc', help="minimum percentage of matched bins required to add contigs to highest-quality bin",
                                 type=float, default=50.0)
    consensus_group.add_argument('-m', '--add_matches',
                                 help="minimum number of matched bins required to 'add' contigs",
                                 type=int, default=3)

    optional_group = consensus_parser.add_argument_group(
        title='optional arguments')
    optional_group.add_argument('--marker_dir',
                                help="directory containing Pfam and TIGRfam marker genes data")
    optional_group.add_argument('--report_min_quality',
                                help="minimum quality of bin to report",
                                type=float, default=10)
    optional_group.add_argument('--simple_headers',
                                help="do not add additional information to FASTA headers",
                                action='store_true')
    optional_group.add_argument('-p', '--bin_prefix',
                                help="prefix for output bins",
                                default='bin')
    optional_group.add_argument('--silent',
                                help="suppress output of logger",
                                action='store_true')
    optional_group.add_argument('-h', '--help',
                                help="show help message",
                                action="help")

    # greedy bin selection
    greedy_parser = subparsers.add_parser('greedy',
                                          formatter_class=CustomHelpFormatter,
                                          description='Greedy bin selection across multiple binning methods.',
                                          add_help=False)

    greedy_parser.add_argument('profile_dir',
                               help="directory with bin profiles (output of 'profile' command)")
    greedy_parser.add_argument('output_dir',
                               help="output directory")

    required_group = greedy_parser.add_argument_group(
        title='required arguments (one of)')
    bin_group = required_group.add_mutually_exclusive_group(required=True)
    bin_group.add_argument('-b', '--bin_dirs',
                           help="directories with bins from different binning methods",
                           nargs='*')
    bin_group.add_argument('-f', '--bin_file',
                           help="tab-separated file indicating directories with bins from binning methods (two columns: method name and directory)")

    greedy_group = greedy_parser.add_argument_group(
        title="optional bin selection arguments")
    greedy_group.add_argument('-w', '--weight',
                              help="weight given to contamination for assessing genome quality",
                              type=float, default=2)
    greedy_group.add_argument('-q', '--sel_min_quality',
                              help="minimum quality of bin to consider during bin selection process",
                              type=float, default=50)
    greedy_group.add_argument('-x', '--sel_min_comp',
                              help="minimum completeness of bin to consider during bin selection process",
                              type=float, default=50)
    greedy_group.add_argument('-y', '--sel_max_cont',
                              help="maximum contamination of bin to consider during bin selection process",
                              type=float, default=10)

    optional_group = greedy_parser.add_argument_group(
        title='optional arguments')
    optional_group.add_argument('--marker_dir',
                                help="directory containing Pfam and TIGRfam marker genes data")
    optional_group.add_argument('--report_min_quality',
                                help="minimum quality of bin to report",
                                type=float, default=10)
    optional_group.add_argument('--simple_headers',
                                help="do not add additional information to FASTA headers",
                                action='store_true')
    optional_group.add_argument('-p', '--bin_prefix',
                                help="prefix for output bins",
                                default='bin')
    optional_group.add_argument('--silent',
                                help="suppress output of logger",
                                action='store_true')
    optional_group.add_argument('-h', '--help',

                                help="show help message",
                                action="help", )

    # unanimous bin selection
    unanimous_parser = subparsers.add_parser('unanimous',
                                             formatter_class=CustomHelpFormatter,
                                             description='Unanimous bin filtering across multiple binning methods.',
                                             add_help=False)

    unanimous_parser.add_argument('profile_dir',
                                  help="directory with bin profiles (output of 'profile' command)")
    unanimous_parser.add_argument('output_dir', help="output directory")

    required_group = unanimous_parser.add_argument_group(
        title='required arguments (one of)')
    bin_group = required_group.add_mutually_exclusive_group(required=True)
    bin_group.add_argument('-b', '--bin_dirs',
                           help="directories with bins from different binning methods",
                           nargs='*')
    bin_group.add_argument('-f', '--bin_file',
                           help="tab-separated file indicating directories with bins from binning methods (two columns: method name and directory)")

    unanimous_group = unanimous_parser.add_argument_group(
        title="optional bin selection arguments")
    unanimous_group.add_argument('-w', '--weight',
                                 help="weight given to contamination for assessing genome quality",
                                 type=float, default=2)
    unanimous_group.add_argument('-q', '--sel_min_quality',
                                 help="minimum quality of bin to consider during bin selection process",
                                 type=float, default=50)
    unanimous_group.add_argument('-x', '--sel_min_comp',
                                 help="minimum completeness of bin to consider during bin selection process",
                                 type=float, default=50)
    unanimous_group.add_argument('-y', '--sel_max_cont',
                                 help="maximum contamination of bin to consider during bin selection process",
                                 type=float, default=10)

    optional_group = unanimous_parser.add_argument_group(
        title='optional arguments')
    optional_group.add_argument('--marker_dir',
                                help="directory containing Pfam and TIGRfam marker genes data")
    optional_group.add_argument('--report_min_quality',
                                help="minimum quality of bin to report",
                                type=float,
                                default=10)
    optional_group.add_argument('--simple_headers',
                                help="do not add additional information to FASTA headers",
                                action='store_true')
    optional_group.add_argument('-p', '--bin_prefix',
                                help="prefix for output bins",
                                default='bin')
    optional_group.add_argument('--silent',
                                help="suppress output of logger",
                                action='store_true')
    optional_group.add_argument('-h', '--help',
                                action="help",
                                help="show help message")

    # compare two sets of bins (e.g., from alternative binning methods)
    compare_parser = subparsers.add_parser('compare',
                                           formatter_class=CustomHelpFormatter,
                                           description='Compare bins from two binning methods.')
    compare_parser.add_argument('assembly_file',
                                help="assembled contigs used to generate bins")
    compare_parser.add_argument('bin_dir1',
                                help="directory containing bins from first method")
    compare_parser.add_argument('bin_dir2',
                                help="directory containing bins from second method")
    compare_parser.add_argument('output_file',
                                help="output file showing overlap between bins")

    compare_parser.add_argument('-x', '--extension1',
                                default='fna',
                                help="extension of bins in directory 1")
    compare_parser.add_argument('-y', '--extension2',
                                default='fna',
                                help="extension of bins in directory 2")
    compare_parser.add_argument('--silent',
                                help="suppress output of logger",
                                action='store_true')

    # ensure sequences are unique
    unique_parser = subparsers.add_parser('unique',
                                          formatter_class=CustomHelpFormatter,
                                          description='Ensure sequences are not assigned to multiple bins.')
    unique_parser.add_argument('bin_dir', help="directory containing bins")
    unique_parser.add_argument('-x', '--extension',
                               default='fna',
                               help="extension of bins (all other files in directory are ignored)")
    unique_parser.add_argument('--silent',
                               help="suppress output of logger",
                               action='store_true')

    # get and check options
    args = None
    if(len(sys.argv) == 1 or sys.argv[1] == '-h' or sys.argv == '--help'):
        print_help()
        sys.exit(0)
    elif(len(sys.argv) == 1 or sys.argv[1] == '-v' or sys.argv == '--version'):
        print(f'UniteM v{version()}')
        sys.exit(0)
    else:
        args = parser.parse_args()

    try:
        logger_setup(args.output_dir, "unitem.log",
                     "UniteM", version(), args.silent)
    except:
        logger_setup(None, "unitem.log", "UniteM", version(), args.silent)

    try:
        parser = OptionsParser()
        parser.parse_options(args)
    except SystemExit:
        print("\nControlled exit resulting from an unrecoverable error or warning.")
        raise
    except:
        print("\nUnexpected error:", sys.exc_info()[0])
        raise
