#!/usr/bin/env python
# Fri Sep 18 16:00:03 EDT 2009

# python imports
import sys
import os
import optparse

# rasmus, compbio imports
from rasmus import treelib, util
from compbio import phylo


o = optparse.OptionParser()
o.add_option("-d", "--dir", dest="dir", action="store_true",
             help="identify trees by their directory")
o.add_option("-s", "--stree", dest="stree", metavar="SPECIES_TREE",
             help="species tree in newick format")
o.add_option("-S", "--smap", dest="smap", metavar="GENE_TO_SPECIES_MAP",
             help="gene to species map")
o.add_option("-T", "--treeext", dest="treeext", metavar="TREE_EXT",
             help="gene tree file extension")
o.add_option("-R", "--reconext", dest="reconext", metavar="RECON_EXT",
             help="reconciliation file extension")

conf, args = o.parse_args()

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

# read args
stree = treelib.read_tree(conf.stree)
gene2species = phylo.read_gene2species(conf.smap)


for filename in sys.stdin:
    filename = filename.rstrip("\n")

    if conf.dir:
        treename = os.path.basename(os.path.dirname(filename))
    else:
        treename = filename

    tree = treelib.read_tree(filename)

    # get recon and events
    if conf.treeext and conf.reconext:
        try:
            reconfile = util.replace_ext(filename, conf.treeext, conf.reconext)
            recon, events = phylo.read_recon_events(reconfile, tree, stree)
        except:
            print >>sys.stderr, "read error: ", filename
            continue
    else:
        # use MPR to build reconciliation and events
        recon = phylo.reconcile(tree, stree, gene2species)
        events = phylo.label_events(tree, recon)

    try:
        dupcons = phylo.dup_consistency(tree, recon, events)
    except:
        print >>sys.stderr, "error: ", filename

    scores = [(node, score) for (node, score) in dupcons.items()]
    scores.sort()

    for node, score in scores:
        print "\t".join([treename, str(node.name),
                         str(recon[node].name), str(score)])

