#!/usr/local/bin/python3.8
###############################################################################
#                                                                             #
#    groopm                                                                   #
#                                                                             #
#    Entry point. See groopm/groopm.py for internals                          #
#                                                                             #
#    Copyright (C) Michael Imelfort                                           #
#                                                                             #
###############################################################################
#                                                                             #
#          .d8888b.                                    888b     d888          #
#         d88P  Y88b                                   8888b   d8888          #
#         888    888                                   88888b.d88888          #
#         888        888d888 .d88b.   .d88b.  88888b.  888Y88888P888          #
#         888  88888 888P"  d88""88b d88""88b 888 "88b 888 Y888P 888          #
#         888    888 888    888  888 888  888 888  888 888  Y8P  888          #
#         Y88b  d88P 888    Y88..88P Y88..88P 888 d88P 888   "   888          #
#          "Y8888P88 888     "Y88P"   "Y88P"  88888P"  888       888          #
#                                             888                             #
#                                             888                             #
#                                             888                             #
#                                                                             #
###############################################################################
#                                                                             #
#    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__ = "Michael Imelfort"
__copyright__ = "Copyright 2012-2014"
__credits__ = ["Michael Imelfort"]
__license__ = "GPL3"
__version__ = "0.3.4"
__maintainer__ = "Michael Imelfort"
__email__ = "mike@mikeimelfort.com"
__status__ = "Released"

###############################################################################

import argparse
import sys
import re
from groopm import groopm

###############################################################################
###############################################################################
###############################################################################
###############################################################################

def printHelp():
    print '''\

                             ...::: GroopM :::...

                     Automagical metagenomic binning FTW!

   -------------------------------------------------------------------------
                                  version: %s
   -------------------------------------------------------------------------

    Typical workflow:

    groopm parse        -> Load the raw data and save to disk
    groopm core         -> Create core bins
    groopm refine       -> Refine these cores a little
    groopm recruit      -> Add more contigs to the cores
    groopm extract      -> Extract binned contigs or reads

    Extra features:

        Utilities:

    groopm merge        -> Merge two or more bins
    groopm split        -> Split a bin into N parts
    groopm delete       -> Delete a bin

        Printing, plotting:

    groopm explore      -> Methods for viewing bin layouts
    groopm plot         -> Plot bins
    groopm highlight    -> Highlight individual bins and apply labels
    groopm flyover      -> Create a movie of your data
    groopm print        -> Print summary statistics

        Import, export:

    groopm dump         -> Write database fields to csv

    USE: groopm OPTION -h to see detailed options
    ''' % __version__

#    groopm import       -> Import data from csv

if __name__ == '__main__':

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

    ##################################################
    # Typical workflow
    ##################################################

    #-------------------------------------------------
    # parse raw data and save
    file_parser = subparsers.add_parser('parse',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='parse raw data and save to disk',
                                        description='Parse raw data and save to disk')
    file_parser.add_argument('dbname', help="name of the database being created")
    file_parser.add_argument('reference', help="fasta file containing bam reference sequences")
    file_parser.add_argument('bamfiles', nargs='+', help="bam files to parse")
    file_parser.add_argument('-t', '--threads', type=int, default=1, help="number of threads to use during BAM parsing")
    file_parser.add_argument('-f', '--force', action="store_true", default=False, help="overwrite existing DB file without prompting")
    file_parser.add_argument('-c', '--cutoff', type=int, default=500, help="cutoff contig size during parsing")

    #-------------------------------------------------
    # load saved data and make bin cores
    core_builder = subparsers.add_parser('core',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='load saved data and make bin cores',
                                        description='Load saved data and make bin cores')
    core_builder.add_argument('dbname', help="name of the database to open")
    core_builder.add_argument('-c', '--cutoff', type=int, default=1500, help="cutoff contig size for core creation")
    core_builder.add_argument('-s', '--size', type=int, default=10, help="minimum number of contigs which define a core")
    core_builder.add_argument('-b', '--bp', type=int, default=1000000, help="cumulative size of contigs which define a core regardless of number of contigs")
    core_builder.add_argument('-f', '--force', action="store_true", default=False, help="overwrite existing DB file without prompting")
    core_builder.add_argument('-g', '--graphfile', help="output graph of micro bin mergers")
    core_builder.add_argument('-p', '--plot', action="store_true", default=False, help="create plots of bins after basic refinement")
    core_builder.add_argument('-m', '--multiplot', default=0, help="create plots during core creation - (0-3) MAKES MANY IMAGES!")

    #-------------------------------------------------
    # refine bins
    bin_refiner = subparsers.add_parser('refine',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='merge similar bins / split chimeric ones',
                                        description='Merge similar bins and split chimeric ones')
    bin_refiner.add_argument('dbname', help="name of the database to open")
#    bin_refiner.add_argument('-b', '--bids', nargs='+', type=int, default=None, help="bin ids to use (None for all)")
    bin_refiner.add_argument('-a', '--auto', action="store_true", default=False, help="automatically refine bins")
    bin_refiner.add_argument('-r', '--no_transform', action="store_true", default=False, help="skip data transformation (3 stoits only)")
    bin_refiner.add_argument('-p', '--plot', action="store_true", default=False, help="create plots of bins after refinement")

    #-------------------------------------------------
    # enlarge bins
    bin_expander = subparsers.add_parser('recruit',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='load saved data and enlarge bins',
                                        description='Recruit more contigs into existing bins')
    bin_expander.add_argument('dbname', help="name of the database to open")
    bin_expander.add_argument('-c', '--cutoff', type=int, default=500, help="cutoff contig size")
    bin_expander.add_argument('-f', '--force', action="store_true", default=False, help="overwrite existing db file without prompting")
    bin_expander.add_argument('-s', '--step', default=200, type=int, help="step size for iterative recruitment")
    bin_expander.add_argument('-i', '--inclusivity', default=2.5, type=float, help="make recruitment more or less inclusive")

    #-------------------------------------------------
    # extract reads and contigs from saved
    bin_extractor = subparsers.add_parser('extract',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='extract contigs or reads based on bin affiliations',
                                        description='Extract contigs or reads based on bin affiliations')
    bin_extractor.add_argument('dbname', help="name of the database to open")
    bin_extractor.add_argument('data', nargs='+', help="data file(s) to extract from, bam or fasta")
    bin_extractor.add_argument('-b', '--bids', nargs='+', type=int, default=None, help="bin ids to use (None for all)")
    bin_extractor.add_argument('-m', '--mode', default="contigs", help="what to extract [reads, contigs]")
    bin_extractor.add_argument('-o', '--out_folder', default="", help="write to this folder (None for current dir)")
    bin_extractor.add_argument('-p', '--prefix', default="", help="prefix to apply to output files")

    bin_extractor.add_argument('-c', '--cutoff', type=int, default=0, help=">>CONTIG MODE ONLY<< cutoff contig size (0 for no cutoff)")

    bin_extractor.add_argument('--mix_bams', action="store_true", default=False, help=">>READ MODE ONLY<< use the same file for multiple bam files")
    bin_extractor.add_argument('--mix_groups', action="store_true", default=False, help=">>READ MODE ONLY<< use the same files for multiple group groups")
    bin_extractor.add_argument('--mix_reads', action="store_true", default=False, help=">>READ MODE ONLY<< use the same files for paired/unpaired reads")
    bin_extractor.add_argument('--interleave', action="store_true", default=False, help=">>READ MODE ONLY<< interleave paired reads in ouput files")
    bin_extractor.add_argument('--headers_only', action="store_true", default=False, help=">>READ MODE ONLY<< extract only (unique) headers")
    bin_extractor.add_argument('--no_gzip', action="store_true", default=False, help="do not gzip output files")

    bin_extractor.add_argument('--mapping_quality', type=int, default=0, help=">>READ MODE ONLY<< mapping quality threshold")
    bin_extractor.add_argument('--use_secondary', action="store_true", default=False, help=">>READ MODE ONLY<< use reads marked with the secondary flag")
    bin_extractor.add_argument('--use_supplementary', action="store_true", default=False, help=">>READ MODE ONLY<< use reads marked with the supplementary flag")
    bin_extractor.add_argument('--max_distance', type=int, default=1000, help=">>READ MODE ONLY<< maximum allowable edit distance from query to reference")

    bin_extractor.add_argument('-v', '--verbose', action="store_true", default=False, help=">>READ MODE ONLY<< be verbose")
    bin_extractor.add_argument('-t', '--threads', type=int, default=1, help=">>READ MODE ONLY<< maximum number of threads to use")

    ##################################################
    # Utilities
    ##################################################

    #-------------------------------------------------
    # combine two or more bins into one
    bin_merger = subparsers.add_parser('merge',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='merge 2 or more bins')
    bin_merger.add_argument('dbname', help="name of the database to open")
    bin_merger.add_argument('bids', nargs='+', type=int, help="bin ids to merge.")
    bin_merger.add_argument('-f', '--force', action="store_true", default=False, help="merge without prompting")

    #-------------------------------------------------
    # split a bin into two parts
    bin_splitter = subparsers.add_parser('split',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='split a bin into n pieces')
    bin_splitter.add_argument('dbname', help="name of the database to open")
    bin_splitter.add_argument('bid', type=int, help="bin id to split")
    bin_splitter.add_argument('parts', type=int, help="number of parts to split the bin into")
    bin_splitter.add_argument('-m', '--mode', default="kmer", help="profile to split on [kmer, cov]")
    bin_splitter.add_argument('-f', '--force', action="store_true", default=False, help="split without prompting")

    #-------------------------------------------------
    # delete bins
    bin_deleter = subparsers.add_parser('delete',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='delete bins')
    bin_deleter.add_argument('dbname', help="name of the database to open")
    bin_deleter.add_argument('bids', nargs='+', type=int, help="bin ids to delete")
    bin_deleter.add_argument('-f', '--force', action="store_true", default=False, help="delete without prompting")

    ##################################################
    # Plotting
    ##################################################

    #-------------------------------------------------
    # visualise all bins
    bin_explorer = subparsers.add_parser('explore',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='explore and validate bins')
    bin_explorer.add_argument('dbname', help="name of the database to open")
    bin_explorer.add_argument('-b', '--bids', nargs='+', type=int, default=None, help="bin ids to plot (None for all)")
    bin_explorer.add_argument('-c', '--cutoff', type=int, default=1000, help="cutoff contig size")
    bin_explorer.add_argument('-m', '--mode', default="binids", help="Exploration mode [binpoints, binids, allcontigs, unbinnedcontigs, binnedcontigs, binassignments, compare, sidebyside, together]")
    bin_explorer.add_argument('-r', '--no_transform', action="store_true", default=False, help="skip data transformation (3 stoits only)")
    bin_explorer.add_argument('-k', '--kmers', action="store_true", default=False, help="include kmers in figure [only used when mode == together]")
    bin_explorer.add_argument('-p', '--points', action="store_true", default=False, help="ignore contig lengths when plotting")
    bin_explorer.add_argument('-C', '--cm', default="HSV", help="set colormap [HSV, Accent, Blues, Spectral, Grayscale, Discrete, DiscretePaired]")

    #-------------------------------------------------
    # flyover  --- usually this is basically an easter egg. If you find it then have fun
    bin_pilot = subparsers.add_parser('flyover',
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                      help='create a purdy flyover plot of the bins you made')
    bin_pilot.add_argument('dbname', help="name of the database to open")
    bin_pilot.add_argument('-b', '--bids', nargs='+', type=int, default=None, help="bin ids to concentrate on (None for all)")
    bin_pilot.add_argument('-c', '--cutoff', type=int, default=1000, help="cutoff contig size")
    bin_pilot.add_argument('-p', '--points', action="store_true", default=False, help="ignore contig lengths when plotting")
    bin_pilot.add_argument('-P', '--prefix', default="file", help="prefix to append to start of output files")
    bin_pilot.add_argument('-t', '--title', default="", help="title to add to output images")
    bin_pilot.add_argument('-B', '--colorbar', action="store_true", default=False, help="show the colorbar")
    bin_pilot.add_argument('-f', '--format', default="jpeg", help="file format output images")
    bin_pilot.add_argument('--fps', type=float, default=10, help="frames per second")
    bin_pilot.add_argument('--totalTime', type=float, default=120., help="how long the movie should go for (seconds)")
    bin_pilot.add_argument('--firstFade', type=float, default=0.05, help="what percentage of the movie is devoted to the unbinned contigs")

    #-------------------------------------------------
    # plot a bin/bins
    bin_plotter = subparsers.add_parser('plot',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='plot bins')
    bin_plotter.add_argument('dbname', help="name of the database to open")
    bin_plotter.add_argument('-b', '--bids', nargs='+', type=int, default=None, help="bin ids to plot (None for all)")
    bin_plotter.add_argument('-t', '--tag', default="BIN", help="tag to add to output filename")
    bin_plotter.add_argument('-f', '--folder', default="", help="save plots in folder")
    bin_plotter.add_argument('-p', '--points', action="store_true", default=False, help="ignore contig lengths when plotting")
    bin_plotter.add_argument('-C', '--cm', default="HSV", help="set colormap [HSV, Accent, Blues, Spectral, Grayscale, Discrete, DiscretePaired]")

    #-------------------------------------------------
    # produce fancy image for publications
    bin_highlighter = subparsers.add_parser('highlight',
                                            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                            help='highlight specific bins')
    bin_highlighter.add_argument('dbname', help="name of the database to open")
    bin_highlighter.add_argument('-P', '--place', action="store_true", default=False, help="use this to help work out azimuth/elevation parameters")
    bin_highlighter.add_argument('-L', '--binlabels', default="", help="replace bin IDs with user specified labels (use 'none' to force no labels)")
    bin_highlighter.add_argument('-C', '--contigcolors', default="", help="specify contig colors")
    bin_highlighter.add_argument('-r', '--radius', action="store_true", default=False, help="draw placement radius to help with label moving")
    bin_highlighter.add_argument('-c', '--cutoff', type=int, default=1000, help="cutoff contig size")
    bin_highlighter.add_argument('-e', '--elevation', type=float, default=25.0, help="elevation in printed image")
    bin_highlighter.add_argument('-a', '--azimuth', type=float, default=-45.0, help="azimuth in printed image")
    bin_highlighter.add_argument('-f', '--file', default="gmview", help="name of image file to produce")
    bin_highlighter.add_argument('-t', '--filetype', default="jpg", help="Type of file to produce")
    bin_highlighter.add_argument('-d', '--dpi', default=300, help="Image resolution")
    bin_highlighter.add_argument('-s', '--show', action="store_true", default=False, help="load image in viewer only")
    bin_highlighter.add_argument('-p', '--points', action="store_true", default=False, help="ignore contig lengths when plotting")
    bin_highlighter.add_argument('-b', '--bids', nargs='+', type=int, default=None, help="bin ids to plot (None for all)")

    #-------------------------------------------------
    # print bin information
    bin_printer = subparsers.add_parser('print',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='print bin information')
    bin_printer.add_argument('dbname', help="name of the database to open")
    bin_printer.add_argument('-b', '--bids', nargs='+', type=int, default=None, help="bin ids to print (None for all)")
    bin_printer.add_argument('-o', '--outfile', default="", help="print to file not STDOUT")
    bin_printer.add_argument('-f', '--format', default='bins', help="output format [bins, contigs]")
    bin_printer.add_argument('-u', '--unbinned', action="store_true", default=False, help="print unbinned contig IDs too")

    ##################################################
    # Import Export
    ##################################################

    #-------------------------------------------------
    # dump data to file
    data_dumper = subparsers.add_parser('dump',
                                        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                        help='write database to text file')
    data_dumper.add_argument('dbname', help="name of the database to open")
    data_dumper.add_argument('-f', '--fields', default="names,bins", help="fields to extract: Build a comma separated list from [names, mers, gc, coverage, tcoverage, ncoverage, lengths, bins] or just use 'all']")
    data_dumper.add_argument('-o', '--outfile', default="GMdump.csv", help="write data to this file")
    data_dumper.add_argument('-s', '--separator', default=",", help="data separator")
    data_dumper.add_argument('--no_headers', action="store_true", default=False, help="don't add headers")

    if False:
        #-------------------------------------------------
        # import from file
        data_importer = subparsers.add_parser('import',
                                              formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                              help='import information from ')
        data_importer.add_argument('dbname', help="name of the database to open")
        data_importer.add_argument('infile', help="file with data to import")
        data_importer.add_argument('-t', '--fields', default="bins", help="data type to import. [bins]")
        data_importer.add_argument('-s', '--separator', default=",", help="data separator")
        data_importer.add_argument('--has_headers', action="store_true", default=False, help="file contains headers")

    ##################################################
    # System
    ##################################################

    #-------------------------------------------------
    # get and check options
    args = None
    if(len(sys.argv) == 1):
        printHelp()
        sys.exit(0)
    elif(sys.argv[1] == '-v' or \
         sys.argv[1] == '--v' or \
         sys.argv[1] == '-version' or \
         sys.argv[1] == '--version'):
        print "GroopM: version %s %s %s" % (__version__,
                                            __copyright__,
                                            __author__)
        sys.exit(0)
    elif(sys.argv[1] == '-h' or \
         sys.argv[1] == '--h' or \
         sys.argv[1] == '-help' or \
         sys.argv[1] == '--help'):
        printHelp()
        sys.exit(0)
    else:
        args = parser.parse_args()

    #-------------------------------------------------
    # do what we came here to do
    try:
        GM_parser = groopm.GroopMOptionsParser(__version__)
        if(False):
            import cProfile
            cProfile.run('GM_parser.parseOptions(args)', 'prof')
            ##########################################
            ##########################################
            # Use this in python console!
            #import pstats
            #p = pstats.Stats('prof')
            #p.sort_stats('cumulative').print_stats(10)
            #p.sort_stats('time').print_stats(10)
            ##########################################
            ##########################################
        else:
            GM_parser.parseOptions(args)
    except:
        print "Unexpected error:", sys.exc_info()[0]
        raise

###############################################################################
###############################################################################
###############################################################################
###############################################################################

