#!/usr/bin/env python
# Thu Dec 22 12:05:13 EST 2011

import sys
import optparse

from rasmus import treelib, util
from compbio import phylo
from compbio.vis import transsvg


o = optparse.OptionParser()
o.set_defaults(tree=[],
               scale=20,
               minlen=1,
               maxlen=10000)

o.add_option("-t", "--tree", dest="tree",
             metavar="<newick file>")
o.add_option("-l", "--scale", dest="scale",
             metavar="<scaling>",
             type="float")
o.add_option("-m", "--minlen", dest="minlen",
             metavar="<minimum branch length>",
             type="int")
o.add_option("-M", "--maxlen", dest="maxlen",
             metavar="<maximum branch length>",
             type="int")
o.add_option("-n", "--names", dest="names", metavar="SPECIES_NAMES",
             help="Use a species names file")
o.add_option("-s", "--stree", dest="stree",
             metavar="<species tree>",
             help="species tree (newick format)")
o.add_option("-r", "--recon", metavar="RECON_FILE")
o.add_option("-b", "--brecon", metavar="BRECON_FILE")
o.add_option("-S", "--smap", dest="smap",
             metavar="<gene2species map>",
             help="mapping of gene names to species names")
o.add_option("-o", "--output", metavar="OUTPUT_FILE",
             default="tree.svg")


#=============================================================================


def process_tree(tree, stree, gene2species):

    #==============================
    # process

    if options.hist or options.hashes:
        # count topology
        hashes.append(phylo.hash_tree(tree, gene2species))

    elif options.histsplit:
        # count splits
        for leaf in tree.leaves():
            tree.rename(leaf.name, gene2species(leaf.name))
        splits.extend(phylo.find_splits(tree))

    elif options.dump:
        # dump mode
        dump_tree(tree)

    elif options.events:
        # add event labels
        assert stree is not None and gene2species is not None
        phylo.count_dup_loss_tree(tree, stree, gene2species)

    else:
        # default mode: display tree
        display_tree(tree, options, gene2species=gene2species, stree=stree)


#=============================================================================
# parse options

conf, args = o.parse_args()


gene2species = phylo.read_gene2species(conf.smap)
stree = treelib.read_tree(conf.stree)
tree = treelib.read_tree(conf.tree)

if conf.names:
    snames = dict(util.read_delim(conf.names))
else:
    snames = None

if conf.recon:
    recon, events = phylo.read_recon_events(open(conf.recon), tree, stree)
    brecon = phylo.recon_events2brecon(recon, events)

else:
    recon = phylo.reconcile(tree, stree, gene2species)
    events = phylo.label_events(tree, recon)
    brecon = phylo.recon_events2brecon(recon, events)

phylo.add_implied_spec_nodes_brecon(tree, brecon)

transsvg.draw_tree(tree, brecon, stree, filename=conf.output,
                   snames=snames)
