#!/usr/bin/env python

"""

    Map bootstrap values onto branches of a tree

"""

# python libs
import optparse

# rasmus libs
from rasmus import treelib
from rasmus import util
from compbio import phylo

o = optparse.OptionParser(usage="%prog [options] <tree file> ...")
o.add_option("-o", "--oldext", dest="oldext",
             metavar="<old tree filename extension>",
             default=".tree")
o.add_option("-n", "--newext", dest="newext",
             metavar="<new tree filename extension>",
             default=".bs.tree")
o.add_option("-b", "--bootext", dest="bootext",
             metavar="<bootstrap trees filename extension>",
             default=".boot.trees")
o.add_option("-r", "--rooted", dest="rooted",
             action="store_true",
             help="treat trees as rooted")
options, files = o.parse_args()

for filename in files:

    # read trees
    trees = treelib.read_trees(filename)
    oneline = len(trees) > 1

    # read bootstrap trees
    boottreefile = util.replace_ext(filename, options.oldext, options.bootext)
    boottrees = treelib.read_trees(boottreefile)

    # output file
    outfile = util.replace_ext(filename, options.oldext, options.newext)
    out = util.open_stream(outfile, "w")

    for tree in trees:
        tree = phylo.add_bootstraps(tree, boottrees, rooted=options.rooted)
        tree.write(out, oneline=oneline)

    out.close()
