#!/usr/bin/env python

# Now allows for no rerooting and for trying multiple reroots of original tree.

import optparse

from rasmus import util
from rasmus import treelib

from compbio import phylo

usage = "%prog [options] <tree file> ..."
parser = optparse.OptionParser(usage=usage)
parser.add_option(
    "-a", "--all", dest="all",
    default=False, action="store_true",
    help="if specified, tries all rerootings of the initial tree")
parser.add_option("-r", "--reroot", dest="reroot",
                  default=False, action="store_true",
                  help="if specified, reroots the initial tree")
parser.add_option("-S", "--smap", dest="smap",
                  metavar="<gene2species map>")
parser.add_option("-s", "--stree", dest="stree",
                  metavar="<species tree>")
parser.add_option("-P", "--parext", dest="parext",
                  metavar="<partition file extension>",
                  default=".split")
parser.add_option("-T", "--treeext", dest="treeext",
                  metavar="<tree filename extension>",
                  default=".tree")
parser.add_option("-Y", "--subtreeext", dest="subtreeext",
                  metavar="<subtree filename extension>",
                  default=".sub.tree")
parser.add_option("--fix-ils", dest="fix_ils",
                  default=False, action="store_true")

options, args = parser.parse_args()
if options.all and (not options.reroot):
    parser.error("--all cannot be specified if --reroot is not specified")

gene2species = phylo.read_gene2species(options.smap)
stree = treelib.read_tree(options.stree)


def normBoot(tree):
    if tree.has_data("boot") and len(tree.root.children) == 2:
        if tree.root.children[0].is_leaf() or \
           tree.root.children[1].is_leaf():
            tree.root.children[0].data["boot"] = 0.0
            tree.root.children[1].data["boot"] = 0.0
            tree.root.data["boot"] = 0.0
        else:
            boot = max(tree.root.children[0].data["boot"],
                       tree.root.children[1].data["boot"])
            tree.root.children[0].data["boot"] = boot
            tree.root.children[1].data["boot"] = boot
            tree.root.data["boot"] = 0.0


def splitTree(tree, stree, gene2species, reroot=False):
    """
    - If root is a pre-speciation duplication:
      - recon_root tree
      - If root is still a pre-speciation duplication
        - remove root and split the tree into two separate trees
        - recurse
    """

    recon = phylo.reconcile(tree, stree, gene2species)
    events = phylo.label_events(tree, recon)
    if options.fix_ils:
        dupcons = phylo.dup_consistency(tree, recon, events)
        events = phylo.fix_ils_errors(events, dupcons, newCopy=False)

    if events[tree.root] == "dup" and \
       recon[tree.root] == stree.root:
        if reroot:
            phylo.recon_root(tree, stree, gene2species, newCopy=False)
            recon = phylo.reconcile(tree, stree, gene2species)
            events = phylo.label_events(tree, recon)
            if options.fix_ils:
                dupcons = phylo.dup_consistency(tree, recon, events)
                events = phylo.fix_ils_errors(events, dupcons, newCopy=False)

        if events[tree.root] == "dup" and \
           recon[tree.root] == stree.root:
            subtrees = []

            for child in tree.root.children:
                subtrees.append(treelib.subtree(tree, child))

                # use the max bootstrap for the root branches
                normBoot(subtrees[-1])

            subtrees2 = []
            for subtree in subtrees:
                subtrees2.extend(
                    splitTree(subtree, stree, gene2species, reroot=True)[1])

            # return subtrees
            return True, subtrees2

        # tree only needs rerooting
        normBoot(tree)
        return True, [tree]

    # nothing needs to change
    return False, [tree]


for filename in args:
    print "reading", filename
    tree = treelib.read_tree(filename)

    if options.all:
        # determine ancestors with genes
        etree = phylo.count_dup_loss_trees([tree], stree, gene2species)
        sps = set([n.name for n in etree if n.data['genes'] > 0])

        # try all possible rerootings
        opt = None
        maxct = util.INF
        for name, node in tree.nodes.iteritems():
            if node.is_leaf():
                continue

            # split
            split, subtrees = splitTree(treelib.reroot(tree, name),
                                        stree, gene2species)

            # check if ancestors with genes are the same
            etree2 = phylo.count_dup_loss_trees(subtrees, stree, gene2species)
            sps2 = set([n.name for n in etree2 if n.data['genes'] > 0])
            if sps != sps2:
                continue

            # count events
            ct = sum([
                phylo.count_dup_loss(T, stree,
                                     phylo.reconcile(T, stree, gene2species))
                for T in subtrees])
            if ct < maxct:
                opt = subtrees
                maxct = ct

        # find optimum
        subtrees = opt
        if len(subtrees) > 1 or\
           (phylo.hash_tree(tree) != phylo.hash_tree(subtrees[0])):
            split = True
        else:
            split = False
    else:
        split, subtrees = splitTree(tree, stree, gene2species,
                                    reroot=options.reroot)

    if split:
        partfilename = util.replace_ext(
            filename, options.treeext, options.parext)
        partfile = open(partfilename, "w")

        print "splitting", filename
        print "  ",
        for i, subtree in enumerate(subtrees):
            print (i+1),
            subtreename = util.replace_ext(filename, options.treeext,
                                           "." + str(i+1) + options.subtreeext)
            subtree.write(subtreename)
            partfile.write("\t".join(subtree.leaf_names()) + "\n")
        print
        partfile.close()
