#!/usr/bin/env python
###############################################################################
#                                                                             #
#    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"
__credits__ = ["Michael Imelfort"]
__license__ = "GPL3"
__version__ = "0.2.1"
__maintainer__ = "Michael Imelfort"
__email__ = "mike@mikeimelfort.com"
__status__ = "Alpha"

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

import argparse
import sys
import re
from groopm import groopm

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

def printHelp():
    print '''\
    
               ...::: groopm :::...
                   
       Automagical metagenomic binning FTW!

    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
    
    Extra features:
    
        Utilities:
      
    groopm merge        -> Merge two or bins
    groopm split        -> Split a bin into N parts
    groopm delete       -> Delete a bin
    
        Printing, plotting:
        
    groopm explore      -> Methods for viewing bin layouts
    groopm print        -> Print bin statistics
    groopm plot         -> Plot bins
    
    USE: groopm OPTION -h to see detailed options
    '''

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('-d', '--dump', action="store_true", default=False, help="dump the contents of all tables to screen")
    file_parser.add_argument('-f', '--force', action="store_true", default=False, help="overwrite existing DB file without prompting")
        
    #-------------------------------------------------
    # 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=3000, help="cutoff contig size for core creation")
    core_builder.add_argument('-s', '--size', type=int, default=5, 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('-p', '--plot', action="store_true", default=False, help="create plots during core creation - 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('-m', '--mode', default='plot', help="refinement mode plot|shuffle")
    bin_refiner.add_argument('-r', '--no_transform', action="store_true", default=False, help="skip data transformation (3 stoits only)")

    #-------------------------------------------------
    # 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 and reads based on bin affiliations',
                                        description='Extract contigs and 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, help="bin ids to use (None for all)")
    bin_extractor.add_argument('-c', '--cutoff', type=int, default=0, help="cutoff contig size, ignored for reads (0 for no cutoff)")
    bin_extractor.add_argument('-m', '--mode', default="contigs", help="what to extract [reads, contigs]")
    bin_extractor.add_argument('-o', '--outfolder', default="", help="write to this folder (None for current dir)")
    bin_extractor.add_argument('-B', '--no_separate_bams', action="store_true", default=False, help="use one file for all stoits")

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

    #-------------------------------------------------
    # load bins and visualise
    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, help="bin ids to plot (None for all)")
    bin_explorer.add_argument('-c', '--cutoff', default=10000, help="cutoff contig size")
    bin_explorer.add_argument('-m', '--mode', default="points", help="Exploration mode [ids, points, profile, compare, flyover, unbinned]")

    #-------------------------------------------------
    # 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, 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='summary', help="output format [full, minimal, summary]")
    bin_printer.add_argument('-u', '--unbinned', action="store_true", default=False, help="print unbinned contig IDs too")

    #-------------------------------------------------
    # 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, help="bin ids to plot (None for all)")
    bin_plotter.add_argument('-s', '--sidebyside', action="store_true", default=False, help="plot bins in 3d side by side instead of to file")
    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")    
        
    #-------------------------------------------------
    # get and check options
    args = None
    if(len(sys.argv) == 1 or sys.argv[1] == '-h' or sys.argv == '--help'):
        printHelp()
        sys.exit(0)
    else:
        args = parser.parse_args()

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