#!/usr/bin/env python

# python libraries
import os, sys, optparse
import time
import random
import StringIO

# dlcpar libraries
import dlcpar
from dlcpar import common
from dlcpar import reconlib
import dlcpar.recon

# rasmus and compbio libraries
from rasmus import treelib, util
from compbio import phylo 

#==========================================================
# parser

VERSION = dlcpar.PROGRAM_VERSION_TEXT

def parse_args():
    """parse input arguments"""
    
    parser = optparse.OptionParser(
        usage = "usage: %prog [options] <gene tree> ...",
        
        version = "%prog " + VERSION,

        description =
             "%prog is a phylogenetic program for finding " +
             "the most parsimonious gene tree-species tree reconciliation " +
	     "by inferring speciation, duplication, loss, and deep coalescence events. " +
	     "See http://compbio.mit.edu/dlcpar for details.",
             
        epilog =
             "Written by Yi-Chieh Wu (yjw@mit.edu), Massachusetts Institute of Technology. " +
             "(c) 2012. Released under the terms of the GNU General Public License.")

    grp_io = optparse.OptionGroup(parser, "Input/Output")
    grp_io.add_option("-s", "--stree", dest="stree",
	              metavar="<species tree>",
		      help="species tree file in newick format")
    grp_io.add_option("-S", "--smap", dest="smap",
	              metavar="<species map>",
	              help="gene to species map")    
    parser.add_option_group(grp_io)

    grp_ext = optparse.OptionGroup(parser, "File Extensions")
    grp_ext.add_option("-I","--inputext", dest="inext",
	               metavar="<input file extension>",
	               default="",
	               help="input file extension (default: \"\")")
    grp_ext.add_option("-O", "--outputext", dest="outext",
                       metavar="<output file extension>",
                       default=".dlcpar",
                       help="output file extension (default: \".dlcpar\")")
    parser.add_option_group(grp_ext)

    grp_costs = optparse.OptionGroup(parser, "Costs")
    grp_costs.add_option("-D", "--dupcost", dest="dupcost",
                         metavar="<dup cost>",
                         default=1.0, type="float",
                         help="duplication cost (default: 1.0)")
    grp_costs.add_option("-L", "--losscost", dest="losscost",
                         metavar="<loss cost>",
                         default=1.0, type="float",
                         help="loss cost (default: 1.0)")
    grp_costs.add_option("-C", "--coalcost", dest="coalcost",
                         metavar="<coal cost>",
                         default=0.5, type="float",
                         help="deep coalescence cost (default: 0.5)")
    grp_costs.add_option("--explicit", dest="explicit",
                         default=False, action="store_true",
                         help="set to ignore extra lineages at implied speciation nodes")
    parser.add_option_group(grp_costs)
    
    grp_heur = optparse.OptionGroup(parser, "Heuristics")
##    grp_heur.add_option("--delay", dest="delay",
##                        default=False, action="store_true",
##                        help="set to allow delay between speciation and coalescence (UNTESTED)")
    parser.set_defaults(delay=False)
    grp_heur.add_option("--no_prescreen", dest="prescreen",
                        default=True, action="store_false",
			help="set to disable prescreen of locus maps")
    grp_heur.add_option("--prescreen_min", dest="prescreen_min",
                        metavar="<prescreen min>",
                        default=5, type="float",
                        help="prescreen maps if min (forward) cost exceeds this value")
    grp_heur.add_option("--prescreen_factor", dest="prescreen_factor",
                        metavar="<prescreen factor>",
                        default=2, type="float",
                        help="prescreen maps with (forward) cost exceeding this factor * min (forward) cost")
    grp_heur.add_option("--max_loci", dest="max_loci",
                        metavar="<max # of loci>",
                        default=-1, type="int",
                        help="maximum # of co-existing loci (in ancestral species), " +\
                             "set to -1 for no limit (default: -1)")
    grp_heur.add_option("--max_dups", dest="max_dups",
                        metavar="<max # of dups>",
                        default=4, type="int",
                        help="maximum # of duplications (in ancestral species), " +\
                             "set to -1 for no limit (default: 4)")
    grp_heur.add_option("--max_losses", dest="max_losses",
                        metavar="<max # of losses>",
                        default=4, type="int",
                        help="maximum # of losses (in ancestral species), " +\
                             "set to -1 for no limit (default: 4)")
    parser.add_option_group(grp_heur)

    grp_misc = optparse.OptionGroup(parser, "Miscellaneous")
    grp_misc.add_option("-x", "--seed", dest="seed",
                        metavar="<random seed>",
                        type="int", default=None,
                        help="random number seed")
    grp_misc.add_option("--output_format", dest="output_format",
                        choices=["dlcpar","dlcoal"], default="dlcpar",
                        metavar="[dlcpar|dlcoal]",
                        help="specify output format (default: dlcpar)")
    parser.add_option_group(grp_misc)

    grp_info = optparse.OptionGroup(parser, "Information")
    common.move_option(parser, "--version", grp_info)
    common.move_option(parser, "--help", grp_info)
    grp_info.add_option("-l", "--log", dest="log",
                        action="store_true",
                        help="if given, output debugging log")
    parser.add_option_group(grp_info)

    options, treefiles = parser.parse_args()
    if len(treefiles) == 0:
        parser.print_help()
        sys.exit(1)

    #=============================
    # check arguments

    # required options
    if (not options.stree) or (not options.smap):
        parser.error("-s/--stree and -S/--smap required")

    # non-negative costs
    if options.dupcost < 0:
        parser.error("-D/--dupcost must be non-negative")
    if options.losscost < 0:
        parser.error("-L/--losscost must be non-negative")
    if options.coalcost < 0:
        parser.error("-C/--coalcost must be non-negative")        

    # prescreen options
    if options.prescreen_min <= 0:
        parser.error("--prescreen_min must be > 0")
    if options.prescreen_factor <= 0:
        parser.error("--prescreen_factor must be > 0")

    # max loci/dups/losses options
    if options.max_loci == -1:
        options.max_loci = util.INF
    elif options.max_loci < 1:
        parser.error("--max_loci must be > 0")

    if options.max_dups == -1:
        options.max_dups = util.INF
    elif options.max_dups < 1:
        parser.error("--max_dups must be > 0")

    if options.max_losses == -1:
        options.max_losses = util.INF
    elif options.max_losses < 1:
        parser.error("--max_losses must be > 0")

    return options, treefiles

      
#==========================================================
# main

def main():
    """main"""
    
    # parse arguments
    options, treefiles = parse_args()

    # read species tree and species map
    stree = treelib.read_tree(options.stree)
    common.check_tree(stree, options.stree)
    gene2species = phylo.read_gene2species(options.smap)

    # process genes trees
    for treefile in treefiles:
        # start logging
        if options.log:
            log_out = util.open_stream(util.replace_ext(treefile, options.inext, options.outext + ".log"), 'w')
        else:
            log_out = common.NullLog()

        # log command
        log_out.write("DLCPar version: %s\n" % VERSION)
        log_out.write("DLCPar executed with the following arguments:\n")
        log_out.write("%s %s\n\n" % (os.path.basename(sys.argv[0]),
                                     ' '.join(map(lambda x: x if x.find(' ') == -1 else "\"%s\"" % x,
                                                  sys.argv[1:]))))
        
        # set random seed
        if options.seed is None:
            options.seed = int(time.time() * 100)
        random.seed(options.seed)
        log_out.write("seed: %d\n" % options.seed)
        log_out.write("\n")

        # read and prepare coal tree
        coal_trees = list(treelib.iter_trees(treefile))

        # multiple coal trees not supported
        if len(coal_trees) > 1:
            raise Exception("unsupported: multiple coal trees per file")

        for coal_tree in coal_trees:
            common.check_tree(coal_tree, treefile)

            # remove bootstrap and distances if they exist
            coal_tree_top = coal_tree.copy()
            for node in coal_tree_top:
                if "boot" in node.data:
                    del node.data["boot"]
                node.dist = 0
            coal_tree_top.default_data.clear()

            # perform reconciliation
            gene_tree, labeled_recon = dlcpar.recon.dlc_recon(
                coal_tree_top, stree, gene2species,
                dupcost=options.dupcost, losscost=options.losscost, coalcost=options.coalcost,
                implied=not options.explicit, delay=options.delay,
                prescreen=options.prescreen, prescreen_min=options.prescreen_min, prescreen_factor=options.prescreen_factor,
                max_loci=options.max_loci, max_dups=options.max_dups, max_losses=options.max_losses,
                log=log_out)

        # write outputs
        out = util.replace_ext(treefile, options.inext, options.outext)
        if options.output_format == "dlcoal":
            coal_tree, recon = reconlib.labeledrecon_to_recon(gene_tree, labeled_recon, stree)
            recon.write(out, coal_tree)
        else:
            labeled_recon.write(out, gene_tree)
        
        # end logging
        if options.log:
            log_out.close()


# main function
if __name__ == "__main__":
    main()
