#!/usr/bin/env python

# use three-tree model as in DLCoal
# but use parsimony cost rather than probabilistic search

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

# dlcpar libraries
import dlcpar
from dlcpar import common
import dlcpar.simplerecon

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

#==========================================================
# 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 (heuristically) 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_search = optparse.OptionGroup(parser, "Search")
    grp_search.add_option("-i", "--iter", dest="iter",
                          metavar="<# iterations>",
                          type="int", default=10,
                          help="number of search iterations (default: 10)")
    grp_search.add_option("", "--nprescreen", dest="nprescreen",
                          metavar="<# prescreens>",
                          type="int", default=20,
                          help="number of prescreening iterations (default: 20)")
    grp_search.add_option("", "--init-locus-tree", dest="init_locus_tree",
                          metavar="<tree file>",
                          help="initial locus tree for search")
    grp_search.add_option("-x", "--seed", dest="seed",
                          metavar="<random seed>",
                          type="int", default=None,
                          help="random number seed")
    parser.add_option_group(grp_search)

    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()

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

    # input gene tree files
    if len(treefiles) == 0:
        parser.print_help()
        sys.exit(1)

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

    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)

    # initial locus tree
    if options.init_locus_tree:
        init_locus_tree = treelib.read_tree(options.init_locus_tree)
    else:
        init_locus_tree = None

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

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

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

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

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

            # perform reconciliation
            maxrecon = dlcpar.simplerecon.dlc_recon(
                coal_tree, stree, gene2species,
                dupcost=options.dupcost, losscost=options.losscost, coalcost=options.coalcost, implied=not options.explicit,
                nsearch=options.iter, nprescreen=options.nprescreen,
                init_locus_tree=init_locus_tree,
                log=log_out)
            locus_trees.append(maxrecon["locus_tree"])

        # make "consensus" reconciliation if multiple coal trees given
        if len(coal_trees) > 1:
            # make consensus locus tree
            coal_tree = phylo.consensus_majority_rule(coal_trees, rooted=True)
            phylo.ensure_binary_tree(coal_tree)
            locus_tree = phylo.consensus_majority_rule(locus_trees, rooted=True)
            phylo.ensure_binary_tree(locus_tree)
            locus_recon = phylo.reconcile(locus_tree, stree, gene2species)
            maxrecon = {
                "coal_recon": phylo.reconcile(coal_tree, locus_tree, lambda x:x),
                "locus_tree": locus_tree,
                "locus_recon": locus_recon,
                "locus_events": phylo.label_events(locus_tree, locus_recon)}
            maxrecon["daughters"] = dlcpar.simplerecon.propose_daughters(
                coal_tree, maxrecon["coal_recon"], locus_tree,
                maxrecon["locus_events"])

        # write outputs
        out = util.replace_ext(treefile, options.inext, options.outext)
        phyloDLC.write_dlcoal_recon(out, coal_tree, maxrecon)

        # end logging
        if options.log:
            log_out.close()

# main function
if __name__ == "__main__":
    sys.exit(main())
