"""
treelib.py
Tree data structures

Contains special features for representing phylogeny.
See compbio.phylo for more.
"""

# python libraries
import copy
import sys
import StringIO

# rasmus libraries
from rasmus import util
from rasmus import textdraw


#============================================================================
# Tree data structures

class TreeNode(object):
    """A class for nodes in a rooted Tree

    Contains fields for branch length 'dist' and custom data 'data'
    """

    def __init__(self, name=None):
        self.name = name
        self.children = []
        self.parent = None
        self.dist = 0
        self.data = {}


    def __iter__(self):
        """Iterate through child nodes"""
        return iter(self.children)


    def copy(self, parent=None, copy_children=True, copy_data=True):
        """Returns a copy of a TreeNode and all of its children"""

        node = TreeNode(self.name)
        node.name = self.name
        node.dist = self.dist
        node.parent = parent
        if copy_data:
            node.data = copy.copy(self.data)
        if copy_children:
            for child in self.children:
                node.children.append(child.copy(node, copy_data=copy_data))

        return node


    def is_leaf(self):
        """Returns True if the node is a leaf (no children)"""
        return len(self.children) == 0


    def recurse(self, func, *args):
        """Applies a function 'func' to the children of a node"""
        for child in self.children:
            func(child, *args)


    def leaves(self):
        """Returns the leaves beneath the node in traversal order"""
        leaves = []

        def walk(node):
            """helper function"""
            if node.is_leaf():
                leaves.append(node)
            for child in node.children:
                walk(child)
        walk(self)

        return leaves


    def leaf_names(self):
        """Returns the leaf names beneath the node in traversal order"""
        return [leaf.name for leaf in self.leaves()]


    def ancestors(self):
        """Returns the ancestors above the node in traversal order"""
        anc = []

        def walk(node):
            """helper function"""
            if node.parent:
                anc.append(node.parent)
                walk(node.parent)
        walk(self)

        return anc


    def ancestor_names(self):
        """Returns the ancestor names above the node in traversal order"""
        return [x.name for x in self.ancestors()]


    def descendants(self):
        """Returns the descendants beneath the node in traversal order"""
        des = []

        def walk(node):
            """helper function"""
            for child in node.children:
                des.append(child)
                walk(child)
        walk(self)

        return des


    def descendant_names(self):
        """Returns the descendant names beneath the node in traversal order"""
        return [x.name for x in self.descendants()]


    def write_data(self, out):
        """Writes the data of the node to the file stream 'out'"""
        out.write(str(self.dist))


    def __repr__(self):
        """Returns a representation of the node"""
        return "<node %s>" % self.name


class BranchData(object):
    """A class for managing branch specific data for a Tree

    By default, this class implements bootstrap data for TreeNode's.

    To incorporate new kinds of branch data, do the following.  Subclass this
    class (say, MyBranchData).  Create Tree's with
    Tree(branch_data=MyBranchData()).  This will ensure your new branch data
    manager is used for manipulations to the tree.  Any tree's copied from the
    tree (via tree.copy()) will also use the same branch manager.
    """

    def __init__(self):
        pass


    def get_branch_data(self, node):
        """Returns branch specific data from a node"""
        if "boot" in node.data:
            return {"boot": node.data["boot"]}
        return {}


    def set_branch_data(self, node, data):
        """Set the branch specific data from 'data' to node.data"""
        if "boot" in data:
            node.data["boot"] = data["boot"]


    def split_branch_data(self, node):
        """Split a branch's data into two copies"""
        if "boot" in node.data:
            return {"boot": node.data["boot"]}, {"boot": node.data["boot"]}
        return {}, {}


    def merge_branch_data(self, data1, data2):
        """Merges the branch data from two neighboring branches into one"""
        if "boot" in data1 and "boot" in data2:
            assert data1["boot"] == data2["boot"], (data1, data2)
            return {"boot": data1["boot"]}
        return {}


class Tree(object):
    """
    Basic rooted tree

    Well suited for phylogenetic trees
    """

    def __init__(self, nextname=1, branch_data=BranchData(),
                 name=None):
        self.nodes = {}
        self.root = None
        self.nextname = nextname
        self.default_data = {}
        self.data = {}
        self.branch_data = branch_data
        self.name = name


    def __repr__(self):
        """Returns a representation of the tree"""
        return "<tree %s>" % (self.name if self.name is not None else
                              hex(id(self)))


    def copy(self, copy_data=True):
        """Returns a copy of the tree"""
        tree = Tree(nextname=self.nextname, name=self.name)

        # copy structure
        if self.root is not None:
            # copy all nodes
            tree.root = self.root.copy(copy_data=copy_data)

            # set all names
            def walk(node):
                """helper function"""
                tree.nodes[node.name] = node
                for child in node.children:
                    walk(child)
            walk(tree.root)

        # copy extra data
        if copy_data:
            tree.copy_data(self)
            tree.copy_node_data(self)

        return tree


    #=========================================
    # iterators

    def __iter__(self):
        """Iterate through nodes of tree"""
        return self.nodes.itervalues()


    def __len__(self):
        """Returns number of nodes in tree"""
        return len(self.nodes)


    def __getitem__(self, key):
        """Returns node by name"""
        return self.nodes[key]


    def __setitem__(self, key, node):
        """Adds a node to the tree"""
        node.name = key
        self.add(node)


    def __contains__(self, name):
        """Returns True if tree has node with name 'name'"""
        return name in self.nodes


    def preorder(self, node=None, is_leaf=lambda x: x.is_leaf()):
        """Iterate through nodes in pre-order traversal"""

        if node is None:
            node = self.root

        queue = [node]

        while queue: # len(queue) > 0
            node = queue.pop()
            yield node

            if not is_leaf(node):
                for child in reversed(node.children):
                    queue.append(child)


    def postorder(self, node=None, is_leaf=lambda x: x.is_leaf()):
        """Iterate through nodes in post-order traversal"""

        if node is None:
            node = self.root

        stack = [[node, 0]]

        while stack: # len(stack) > 0
            node, i = stack[-1]

            if i < len(node.children) and not is_leaf(node):
                stack.append([node.children[i], 0])
                stack[-2][1] += 1
            else:
                yield node
                stack.pop()


    def inorder(self, node=None, is_leaf=lambda x: x.is_leaf()):
        """Iterate through nodes with in-order traversal"""

        if node is None:
            node = self.root

        stack = [[node, 0]]

        while stack: # len(stack) > 0
            node, i = stack[-1]

            if node.is_leaf():
                yield node
                stack.pop()

            elif i < len(node.children) and not is_leaf(node):
                assert len(node.children) == 2, node.name

                if i == 1:
                    # left has been visited
                    # yield current node then visit right
                    yield node

                # recurse into children
                stack.append([node.children[i], 0])
                stack[-2][1] += 1
            else:
                stack.pop()


    #=============================
    # structure functions

    def make_root(self, name=None):
        """Create a new root node"""
        if name is None:
            name = self.new_name()
        self.root = TreeNode(name)
        return self.add(self.root)


    def add(self, node):
        """Add a node to the tree
           Does not add node to any specific location (use add_child instead).
        """
        self.nodes[node.name] = node
        return node


    def add_child(self, parent, child):
        """
        Add a child node to an existing node 'parent' in the tree
        """
        assert parent != child, (parent.name, child.name)
        self.nodes[child.name] = child
        self.nodes[parent.name] = parent
        child.parent = parent
        parent.children.append(child)
        return child


    def new_node(self, name=None):
        """Add a new node with name 'name' to the tree"""
        if name is None:
            name = self.new_name()
        return self.add(TreeNode(name))


    def remove(self, node):
        """
        Removes a node from a tree.
        Notifies parent (if it exists) that node has been removed.
        Updates node.parent to None.
        """

        if node.parent:
            node.parent.children.remove(node)
            node.parent = None
        del self.nodes[node.name]


    def remove_child(self, parent, child):
        """
        Removes a child node from an existing node 'parent' in the tree.
        Updates the parent-child relationships but does NOT remove the nodes
        (use remove instead).
        """

        assert parent != child and child.parent == parent, (
            parent.name, child.name)
        parent.children.remove(child)
        child.parent = None


    def remove_tree(self, node):
        """
        Removes subtree rooted at 'node' from tree.
        Notifies parent (if it exists) that node has been removed.
        Updates node.parent to None.
        """

        def walk(node):
            """helper function"""
            if node.name in self.nodes:
                del self.nodes[node.name]
            for child in node.children:
                walk(child)
        walk(node)

        if node.parent:
            node.parent.children.remove(node)
            node.parent = None


    def rename(self, oldname, newname):
        """Rename a node in the tree"""
        if oldname != newname:
            assert newname not in self.nodes, newname
            node = self.nodes[oldname]
            del self.nodes[oldname]
            self.nodes[newname] = node
            node.name = newname


    def new_name(self):
        """Returns a new node name that should be unique in the tree"""
        name = self.nextname
        self.nextname += 1
        return name


    def unique_name(self, name, names, sep="_"):
        """Create a new unique name not already in names"""
        i = 1
        name2 = name
        while name2 in names:
            name2 = name + sep + str(i)
            i += 1
        names.add(name2)
        return name2


    def add_tree(self, parent, child_tree):
        """Add subtree child_tree to this tree as child of parent."""

        # Merge nodes and change the names of childTree names if they conflict
        # with existing names
        self.merge_names(child_tree)
        self.add_child(parent, child_tree.root)


    def replace_tree(self, node, child_tree):
        """Remove node and replace it with the root of child_tree"""
        node_parent = node.parent
        self.remove_tree(node)
        self.add_tree(node_parent, child_tree)


    def merge_names(self, tree2):
        """Merge the node names from tree2 into this tree.
           Change any names that conflict"""

        for name in tree2.nodes:
            if name in self.nodes:
                name2 = self.new_name()
                self.nodes[name2] = tree2.nodes[name]
                self.nodes[name2].name = name2
            else:
                # make sure I do not issue a name that matches this one
                if isinstance(name, int):
                    if name >= self.nextname:
                        self.nextname = name + 1
                self.nodes[name] = tree2.nodes[name]


    def clear(self):
        """Clear all nodes from tree"""
        self.nodes = {}
        self.root = None


    def leaves(self, node=None):
        """Return the leaves of the tree in order"""
        if node is None:
            node = self.root
            if node is None:
                return []
        return node.leaves()


    def leaf_names(self, node=None):
        """Returns the leaf names of the tree in order"""
        return [leaf.name for leaf in self.leaves(node)]


    def ancestors(self, node):
        """Returns the ancestors in order"""
        return node.ancestors()


    def ancestor_names(self, node):
        """Returns the ancestor names of the tree in order"""
        return node.ancestor_names()


    def descendants(self, node):
        """Returns the descendants in order"""
        return node.descendants()


    def descendant_names(self, node):
        """Returns the descendant names in order"""
        return node.descendant_names()


    #===============================
    # data functions

    def has_data(self, dataname):
        """Does the tree contain 'dataname' in its extra data"""
        return dataname in self.default_data


    def copy_data(self, tree):
        """Copy tree data to another"""
        self.branch_data = tree.branch_data
        self.default_data = copy.copy(tree.default_data)
        self.data = copy.copy(tree.data)


    def copy_node_data(self, tree):
        """Copy node data to another tree"""
        for name, node in self.nodes.iteritems():
            if name in tree.nodes:
                node.data = copy.copy(tree.nodes[name].data)
        self.set_default_data()


    def set_default_data(self):
        """Set default values in each node's data"""
        for node in self.nodes.itervalues():
            for key, val in self.default_data.iteritems():
                node.data.setdefault(key, val)


    def clear_data(self, *keys):
        """Clear tree data"""
        for node in self.nodes.itervalues():
            if len(keys) == 0:
                node.data = {}
            else:
                for key in keys:
                    if key in node.data:
                        del node.data[key]


    #======================================================================
    # branch data functions
    # forward branch data calles to branch data manager

    def get_branch_data(self, node):
        """Returns branch specific data from a node"""
        return self.branch_data.get_branch_data(node)


    def set_branch_data(self, node, data):
        """Set the branch specific data from 'data' to node.data"""
        return self.branch_data.set_branch_data(node, data)


    def split_branch_data(self, node):
        """Split a branch's data into two copies"""
        return self.branch_data.split_branch_data(node)


    def merge_branch_data(self, data1, data2):
        """Merges the branch data from two neighboring branches into one"""
        return self.branch_data.merge_branch_data(data1, data2)


    #=======================================================================
    # input and output

    def read_data(self, node, data, namefunc=lambda name: name):
        """Default data reader: reads optional bootstrap and branch length"""

        # also parse nhx comments
        data = read_nhx_data(node, data)

        if ":" in data:
            boot, dist = data.split(":")
            node.dist = float(dist)

            if boot: # len(boot) > 0
                if boot.isdigit():
                    node.data["boot"] = int(boot)
                else:
                    try:
                        node.data["boot"] = float(boot)
                    except ValueError:
                        # treat as node name
                        name = boot.strip()
                        if name and isinstance(node.name, int):
                            node.name = namefunc(name)
                        if name and node.name is None:
                            node.name = name
        else:
            data = data.strip()

            # treat as name
            if data:
                node.name = namefunc(data)


    def write_data(self, node, write_dist=True,
                   namefunc=lambda name: name):
        """Default data writer: writes optional bootstrap and branch length"""

        string = ""
        if "boot" in node.data and \
           not node.is_leaf() and \
           self.root != node:
            if isinstance(node.data["boot"], int):
                string += "%d" % node.data["boot"]
            else:
                string += "%f" % node.data["boot"]
        else:
            # see if internal node names exist
            if not node.is_leaf() and isinstance(node.name, str):
                string += namefunc(node.name)

        if node.dist != 0 or write_dist:
            string += ":%f" % node.dist
        return string


    def read_newick(self, filename, read_data=None, namefunc=lambda name: name):
        """
        Reads newick tree format from a file stream

        You can specify a specialized node data reader with 'readData'
        """
        read_tree(filename, read_data=read_data, tree=self, namefunc=namefunc)


    def write_newick(self, out=sys.stdout, write_data=None, oneline=False,
                     root_data=True, namefunc=lambda name: name):
        """Write the tree in newick notation"""
        write_newick(self, util.open_stream(out, "w"),
                     write_data=write_data, oneline=oneline,
                     root_data=root_data, namefunc=namefunc)

    write = write_newick


    def get_one_line_newick(self, root_data=True, write_data=None,
                            namefunc=lambda name: name):
        """Get a presentation of the tree in a oneline string newick format"""
        stream = StringIO.StringIO()
        self.write(stream, oneline=True,
                   write_data=write_data, root_data=root_data,
                   namefunc=namefunc)
        return stream.getvalue()


#============================================================================
# Input/Output functions

def read_tree(infile, read_data=None, tree=None, namefunc=lambda name: name):
    """Read a tree from a file stream"""
    infile = util.open_stream(infile)
    return parse_newick(infile, read_data=read_data, tree=tree,
                        namefunc=namefunc)


def read_newick(infile, read_data=None, tree=None, namefunc=lambda name: name):
    """Read a tree from a file stream"""
    infile = util.open_stream(infile)
    return parse_newick(infile, read_data=read_data, tree=tree,
                        namefunc=namefunc)


def iter_trees(treefile, read_data=None, namefunc=lambda name: name):
    """Iterate through multiple trees from a tree file"""

    infile = util.open_stream(treefile)

    # ensure at least one tree in file
    yield read_tree(infile, read_data=read_data, namefunc=namefunc)
    try:
        while True:
            yield read_tree(infile, read_data=read_data, namefunc=namefunc)
    except Exception:
        pass


def read_trees(filename, read_data=None, namefunc=lambda name: name):
    """Read multiple trees from a tree file"""
    return list(iter_trees(filename, read_data=read_data, namefunc=namefunc))


def tokenize_newick(infile):
    """Iterate through the tokens in a stream in newick format

    infile -- a string or file stream
    """

    def iter_stream(infile):
        """helper function"""
        while True:
            yield infile.read(1)

    if not isinstance(infile, basestring):
        infile = iter_stream(infile)
    else:
        infile = iter(infile)

    word = []
    for c in infile:
        if c == "":
            # EOF encountered
            break

        elif c in " \t\n":
            # skip white space
            if word:
                yield "".join(word)
                word[:] = []

        elif c in ";(),:[]":
            # special tokens
            if word:
                yield "".join(word)
                word[:] = []

            if c == "[":
                # parse comment
                word.append(c)
                for c in infile:
                    word.append(c)
                    if c == "]":
                        break
                yield "".join(word)
                word[:] = []
            else:
                yield c
        else:
            # word token
            word.append(c)

    if word:
        yield "".join(word)
        word[:] = []


def parse_newick(infile, read_data=None, tree=None,
                 namefunc=lambda name: name):
    """Parse a newick string or stream

    infile    -- a string or file stream
    read_data -- an optional function for reading node data fields
    tree      -- an optional tree to populate
    namefunc  -- an optional map for node names
    """

    # node stack
    ancestors = []

    # create tree
    if tree is None:
        tree = Tree()
    if read_data is None:
        read_data = tree.read_data

    # create root
    node = TreeNode()
    tree.root = node
    nodes = [node]

    # process token stream
    tokens = tokenize_newick(infile)
    token = None
    data = []
    empty = True
    try:
        while True:
            prev_token = token
            token = tokens.next()
            empty = False

            if token == '(':  # new branchset
                if data:
                    read_data(node, "".join(data), namefunc)
                    data = []
                child = TreeNode()
                nodes.append(child)
                child.parent = node
                node.children.append(child)
                ancestors.append(node)
                node = child

            elif token == ',':  # another branch
                if data:
                    read_data(node, "".join(data), namefunc)
                    data = []
                parent = ancestors[-1]
                child = TreeNode()
                nodes.append(child)

                child.parent = parent
                parent.children.append(child)
                node = child

            elif token == ')':  # optional name next
                if data:
                    read_data(node, "".join(data), namefunc)
                    data = []
                node = ancestors.pop()

            elif token == ':':  # optional length next
                data.append(token)

            elif token == ';':  # end of tree
                if data:
                    read_data(node, "".join(data), namefunc)
                    data = []
                break

            else:
                if prev_token in '(,':
                    node.name = namefunc(token)

                elif prev_token in '):':
                    data.append(token)

                else:
                    data.append(token)

    except StopIteration:
        if empty:
            raise Exception("Empty tree")

    # setup node names
    names = set()
    for node in nodes:
        if node.name is None:
            node.name = tree.new_name()
        node.name = tree.unique_name(node.name, names)
        tree.nodes[node.name] = node

    # test for bootstrap presence
    for node in nodes:
        if "boot" in node.data:
            tree.default_data["boot"] = 0
            break
    tree.set_default_data()

    return tree


def write_newick(tree, out=sys.stdout, write_data=None, oneline=False,
                 root_data=True, namefunc=lambda name: name):
    """Write the tree in newick notation"""
    write_newick_node(tree, tree.root, util.open_stream(out, "w"),
                      write_data=write_data, oneline=oneline,
                      root_data=root_data, namefunc=namefunc)


def write_newick_node(tree, node, out=sys.stdout,
                      depth=0, write_data=None, oneline=False,
                      root_data=True, namefunc=lambda name: name):
    """Write the node in newick format to the out file stream"""

    # default data writer
    if write_data is None:
        write_dist = any(node.dist != 0 for node in tree)
        write_data = lambda node: tree.write_data(node, write_dist=write_dist, namefunc=namefunc)

    if not oneline:
        out.write(" " * depth)

    if not node.children: # len(node.children) == 0:
        # leaf
        out.write(str(namefunc(node.name)))
    else:
        # internal node
        if oneline:
            out.write("(")
        else:
            out.write("(\n")
        for child in node.children[:-1]:
            write_newick_node(tree, child, out, depth+1,
                              write_data=write_data, oneline=oneline,
                              namefunc=namefunc)
            if oneline:
                out.write(",")
            else:
                out.write(",\n")
        write_newick_node(tree, node.children[-1], out, depth+1,
                          write_data=write_data, oneline=oneline,
                          namefunc=namefunc)
        if oneline:
            out.write(")")
        else:
            out.write("\n" + (" " * depth) + ")")

    # don't print data for root node
    if depth == 0:
        if root_data:
            out.write(write_data(node))
        if oneline:
            out.write(";")
        else:
            out.write(";\n")
    else:
        out.write(write_data(node))


#=============================================================================
# NHX format

def parse_nhx_comment(comment):
    """Parse a NHX comment"""
    for pair in comment.split(":"):
        if "=" in pair:
            yield pair.split("=")


def format_nhx_comment(data):
    """Format a NHX comment"""
    return "[&&NHX:" + ":".join("%s=%s" % (k, v)
                                for k, v in data.iteritems()) + "]"


def parse_nhx_data(text):
    """Parse the data field of an NHX file"""
    data = None

    if "[" in text:
        data = {}
        i = text.find("[")
        j = text.find("]")
        comment = text[i+1:j]
        text = text[:i]

        if comment.startswith("&&NHX:"):
            for k, v in parse_nhx_comment(comment[6:]):
                data[k] = v

    return text, data


def read_nhx_data(node, text):
    """Read data function for parsing the data field of an NHX file"""

    text, data = parse_nhx_data(text)
    if data:
        node.data.update(data)
    return text


def write_nhx_data(node):
    """Write data function for writing the data field of an NHX file"""
    text = Tree().write_data(node)

    # write bootstraps separately
    if "boot" in node.data:
        data = node.data.copy()
        del data["boot"]
    else:
        data = node.data

    if data:
        text += format_nhx_comment(data)
    return text


#============================================================================
# Misc. functions for manipulating trees

def assert_tree(tree):
    """Assert that the tree data structure is internally consistent"""

    visited = set()

    def walk(node):
        """helper function"""
        assert node.name in tree.nodes, (tree.name, node.name)
        assert node.name not in visited, (tree.name, node.name)
        visited.add(node.name)
        if node.parent:
            assert node in node.parent.children, (tree.name, node.name)
        for child in node.children:
            assert child.parent == node, (tree.name, node.name, child.name)
        node.recurse(walk)
    walk(tree.root)
    assert tree.root.parent is None, (tree.name, tree.root.name)
    assert len(tree.nodes) == len(visited), (
        "%d %d" % (len(tree.nodes), len(visited)))
    assert tree.root.parent is None
    assert len(tree.nodes) == len(visited), (
        "%d %d" % (len(tree.nodes), len(visited)))


def is_binary(tree):
    """Returns True if tree is binary"""

    rooted = is_rooted(tree)
    for node in tree:
        if not node.is_leaf():
            if (not node.parent) and (not rooted):
                if len(node.children) != 3:
                    return False
            else:
                if len(node.children) != 2:
                    return False
    return True


def lca(nodes):
    """Returns the Least Common Ancestor (LCA) of a list of nodes"""

    if len(nodes) == 1:
        return nodes[0]
    elif len(nodes) > 2:
        return lca([lca(nodes[:2])] + nodes[2:])
    elif len(nodes) == 2:
        node1, node2 = nodes
        set1 = set([node1])
        set2 = set([node2])

        while True:
            if node1 in set2:
                return node1
            if node2 in set1:
                return node2
            if node1.parent is not None:
                node1 = node1.parent
            if node2.parent is not None:
                node2 = node2.parent

            set1.add(node1)
            set2.add(node2)
    else:
        raise Exception("No nodes given")


def find_dist(tree, name1, name2):
    """Returns the branch distance between two nodes in a tree"""

    if name1 not in tree.nodes or name2 not in tree.nodes:
        raise Exception("nodes '%s' and '%s' are not in tree" %
                        (name1, name2))

    # find root path for node1
    node1 = tree.nodes[name1]
    path1 = [node1]
    while node1 != tree.root:
        node1 = node1.parent
        path1.append(node1)

    # find root path for node2
    node2 = tree.nodes[name2]
    path2 = [node2]
    while node2 != tree.root:
        node2 = node2.parent
        path2.append(node2)

    # find when paths diverge
    i = 1
    while i <= len(path1) and i <= len(path2) and (path1[-i] == path2[-i]):
        i += 1

    dist = 0
    for j in range(i, len(path1)+1):
        dist += path1[-j].dist
    for j in range(i, len(path2)+1):
        dist += path2[-j].dist

    return dist


def descendants(node, lst=None):
    """Return a list of all the descendants beneath a node"""
    if lst is None:
        lst = []
    for child in node.children:
        lst.append(child)
        descendants(child, lst=lst)
    return lst


def count_descendants(node, sizes=None):
    """Returns a dict with number of leaves beneath each node"""
    if sizes is None:
        sizes = {}

    if node.children: # len(node.children) > 0:
        sizes[node] = 0
        for child in node.children:
            count_descendants(child, sizes)
            sizes[node] += sizes[child]
    else:
        sizes[node] = 1

    return sizes


def subtree(tree, node):
    """Return a copy of a subtree of 'tree' rooted at 'node'"""

    # make new tree
    tree2 = Tree(nextname=tree.new_name())

    # copy nodes and data
    tree2.root = node.copy()
    tree2.copy_data(tree)
    tree2.root.parent = None

    # add nodes
    def walk(node):
        """helper function"""
        tree2.add(node)
        node.recurse(walk)
    walk(tree2.root)

    return tree2


def max_disjoint_subtrees(tree, subroots):
    """Returns a list of rooted subtrees with atmost one node from
       the list 'subroots'
    """

    marks = {}

    # mark the path from each subroot to the root
    for subroot in subroots:
        ptr = subroot
        while ptr is not None:
            lst = marks.setdefault(ptr, [])
            lst.append(subroot)
            ptr = ptr.parent

    # subtrees are those trees with nodes that have at most one mark
    subroots2 = []

    def walk(node):
        """helper function"""
        marks.setdefault(node, [])
        if (len(marks[node]) < 2 \
            and (not node.parent or len(marks[node.parent]) >= 2)):
            subroots2.append(node)
        node.recurse(walk)
    walk(tree.root)

    return subroots2


def tree2graph(tree):
    """Convert a tree to a graph data structure (sparse matrix)"""
    mat = {}

    # init all rows of adjacency matrix to
    for name in tree.nodes:
        mat[name] = {}

    for name, node in tree.nodes.iteritems():
        for child in node.children:
            mat[name][child.name] = child.dist

        if node.parent:
            mat[name][node.parent.name] = node.dist

    return mat


def graph2tree(mat, root, closedset=None):
    """Convert a graph to a tree data structure"""

    if closedset is None:
        closedset = set()
    tree = Tree()

    def walk(name):
        """helper function"""
        node = TreeNode(name)
        node.dist = 0
        closedset.add(name)
        for child in mat[name]:
            if child not in closedset:
                child_node = walk(child)
                child_node.dist = mat[name][child]
                tree.add_child(node, child_node)
        return node
    tree.root = walk(root)

    tree.nextname = max(name for name in tree.nodes if isinstance(name, int))

    return tree


def remove_single_children(tree, simplify_root=True):
    """
    Remove all nodes from the tree that have exactly one child

    Branch lengths are added together when node is removed.
    """

    # find single children
    removed = [node
               for node in tree
               if len(node.children) == 1 and node.parent]

    # actually remove children
    for node in removed:
        newnode = node.children[0]

        # add distance
        newnode.dist += node.dist

        # change parent and child pointers
        newnode.parent = node.parent
        index = node.parent.children.index(node)
        node.parent.children[index] = newnode

        # remove old node
        del tree.nodes[node.name]

    # remove singleton from root
    if simplify_root and tree.root and len(tree.root.children) == 1:
        oldroot = tree.root
        tree.root = tree.root.children[0]
        oldroot.children = []
        tree.remove(oldroot)
        tree.root.parent = None
        tree.root.dist += oldroot.dist

    return removed


def remove_exposed_internal_nodes(tree, leaves=None):
    """
    Remove all leaves that were originally internal nodes

    leaves -- a list of original leaves that should stay

    if leaves is not specified, only leaves with strings as names will be kept
    """

    if leaves is not None:
        stay = set(leaves)
    else:
        # use the fact that the leaf name is a string to determine
        # wether to keep it
        stay = set()
        for leaf in tree.leaves():
            if isinstance(leaf.name, basestring):
                stay.add(leaf)

    # post order traverse tree
    def walk(node):
        """helper function"""
        # keep a list of children to visit, since they may remove themselves
        for child in list(node.children):
            walk(child)

        if node.is_leaf() and node not in stay:
            tree.remove(node)
    walk(tree.root)


def subtree_by_leaves(tree, leaves=None, keep_single=False,
                      simplify_root=True):
    """
    Remove any leaf not in leaves set

    leaves        -- a list of leaves that should stay
    keep_single   -- if False, remove all single child nodes
    simplify_root -- if True, basal branch is removed when removing single
                     children nodes
    """

    stay = set(leaves)

    # post order traverse tree
    def walk(node):
        """helper function"""
        # keep a list of children to visit, since they may remove themselves
        for child in list(node.children):
            walk(child)

        if node.is_leaf() and node not in stay:
            tree.remove(node)

    if stay: # len(stay) > 0
        tree.clear()
    else:
        walk(tree.root)

    if not keep_single:
        remove_single_children(tree, simplify_root=simplify_root)

    return tree


def subtree_by_leaf_names(tree, leaf_names, keep_single=False, new_copy=False):
    """Returns a subtree with only the leaves specified"""

    if new_copy:
        tree = tree.copy()
    return subtree_by_leaves(tree, [tree.nodes[x] for x in leaf_names],
                             keep_single=keep_single)


def reorder_tree(tree, tree2, root=True, leafmap=lambda leaf: leaf.name):
    """Reorders the branches of tree to match tree2"""

    if root:
        # reroot tree to match tree2
        root_branches = [set(map(leafmap, n.leaves()))
                         for n in tree2.root.children]

        def walk(node):
            """helper function"""
            if node.is_leaf():
                leaves = set([leafmap(node)])
            else:
                leaves = set()
                for child in node.children:
                    l = walk(child)
                    if l is None:
                        return None
                    leaves = leaves.union(l)

            if leaves in root_branches:
                # root found, terminate search
                reroot(tree, node.name, new_copy=False)
                return None

            return leaves
        walk(tree.root)

    # reorder tree to match tree2
    leaf_lookup = util.list2lookup(map(leafmap, tree2.leaves()))

    def mean(lst):
        """helper function"""
        return sum(lst) / float(len(lst))

    def walk2(node):
        """helper function"""
        if node.is_leaf():
            return set([leafmap(node)])

        # get leaves
        leaf_sets = []
        for child in node.children:
            leaf_sets.append(walk2(child))

        scores = [mean([leaf_lookup[i] for i in l]) for l in leaf_sets]
        rank = util.sort_index(scores)
        node.children = [node.children[i] for i in rank]

        # return union
        ret = leaf_sets[0]
        for l in leaf_sets[1:]:
            ret = ret.union(l)
        return ret

    walk2(tree.root)


def set_tree_topology(tree, tree2):
    """
    Changes the topology of tree to match tree2

    trees must have nodes with the same names
    """
    nodes = tree.nodes
    nodes2 = tree2.nodes

    for node in tree:
        node2 = nodes2[node.name]

        # set parent
        if node2.parent:
            node.parent = nodes[node2.parent.name]
        else:
            node.parent = None

        # set children
        if node.is_leaf():
            assert node2.is_leaf()
        else:
            # copy child structure
            node.children[:] = [nodes[n.name] for n in node2.children]

    tree.root = nodes[tree2.root.name]


#=============================================================================
# Rerooting functions

def is_rooted(tree):
    """Returns True if tree is rooted"""
    return len(tree.root.children) <= 2


def unroot(tree, new_copy=True):
    """Return an unrooted copy of tree"""

    if new_copy:
        tree = tree.copy()

    nodes = tree.root.children
    if len(nodes) == 2 and not (nodes[0].is_leaf() and nodes[1].is_leaf()):
        dist = nodes[0].dist + nodes[1].dist
        data = tree.merge_branch_data(nodes[0].data, nodes[1].data)
        if len(nodes[0].children) < 2:
            nodes.reverse()
        tree.add_child(nodes[0], nodes[1])
        nodes[1].dist = dist
        tree.set_branch_data(nodes[1], data)
        nodes[0].dist = 0
        tree.set_branch_data(nodes[0], {})
        nodes[0].parent = None

        # replace root
        del tree.nodes[tree.root.name]
        tree.root = nodes[0]
    return tree


def reroot(tree, rootname, on_branch=True, new_copy=True, keep_name=False):
    """
    Change the rooting of a tree
    """

    # TODO: remove new_copy (or assert new_copy=False)
    if new_copy:
        tree = tree.copy()

    # handle trivial case
    if ((not on_branch \
         and tree.root.name == rootname) \
         or (on_branch \
             and rootname in [x.name for x in tree.root.children] \
             and len(tree.root.children) == 2)):
        return tree
    assert not on_branch or rootname != tree.root.name, "No branch specified"

    if keep_name:
        assert on_branch  # can only keep name if root is in middle of branch
        oldroot = tree.root.name

    unroot(tree, new_copy=False)

    # handle trivial case
    if not on_branch and tree.root.name == rootname:
        return tree

    if on_branch:
        # add new root in middle of branch
        if keep_name:
            newnode = TreeNode(oldroot)
        else:
            newnode = TreeNode(tree.new_name())
        node1 = tree.nodes[rootname]
        rootdist = node1.dist
        rootdata1, rootdata2 = tree.split_branch_data(node1)
        node1.dist = rootdist / 2.0
        tree.set_branch_data(node1, rootdata1)
        newnode.dist = rootdist / 2.0
        tree.set_branch_data(newnode, rootdata2)

        node2 = node1.parent
        node2.children.remove(node1)
        tree.add_child(newnode, node1)
        tree.add_child(node2, newnode)

        ptr = node2
        ptr2 = newnode
        newroot = newnode
    else:
        # root directly on node
        ptr2 = tree.nodes[rootname]
        ptr = ptr2.parent
        newroot = ptr2

    newroot.parent = None

    # reverse parent child relationship of all nodes on path node1 to root
    next_dist = ptr2.dist
    next_data = tree.get_branch_data(ptr2)
    ptr2.dist = 0
    while True:
        next_ptr = ptr.parent
        ptr.children.remove(ptr2)
        tree.add_child(ptr2, ptr)

        tmp = ptr.dist
        tmp_data = tree.get_branch_data(ptr)
        ptr.dist = next_dist
        tree.set_branch_data(ptr, next_data)
        next_dist = tmp
        next_data = tmp_data

        ptr2 = ptr
        ptr = next_ptr

        if next_ptr is None:
            break
    tree.root = newroot

    return tree


def midpoint_root(tree):
    """Reroot a tree using midpoint rerooting"""

    # get maximum distance from leaves to each node
    depths = {}
    for node in tree.postorder():
        if node.is_leaf():
            depths[node] = (0.0, node)
        else:
            depths[node] = max((c.dist + depths[c][0], depths[c][1])
                               for c in node.children)

    # find maximum path
    dists = []
    for node in tree:
        if node.is_leaf():
            continue
        assert len(node.children) != 1
        tmp = sorted([(c.dist + depths[c][0], depths[c][1], c)
                      for c in node.children])
        dists.append((tmp[-1][0] + tmp[-2][0], node,
                      tmp[-1][2], tmp[-1][1],
                      tmp[-2][2], tmp[-2][1]))

    maxdist, top, child1, leaf1, child2, leaf2 = max(dists)
    middist = maxdist / 2.0

    # find longer part of path
    if depths[child1][0] + child1.dist >= middist:
        ptr = leaf1
    else:
        ptr = leaf2

    # find branch that contains midpoint
    dist = 0.0
    while ptr != top:
        if ptr.dist + dist >= middist:
            # reroot tree
            reroot(tree, ptr.name, on_branch=True, new_copy=False)

            # fixup branch lengths and return
            pdist = sum(c.dist for c in tree.root.children)
            other = [c for c in tree.root.children if c != ptr][0]
            ptr.dist = middist - dist
            other.dist = pdist - ptr.dist
            return tree

        dist += ptr.dist
        ptr = ptr.parent

    raise AssertionError("Could not find branch with midpoint")


#=============================================================================
# ages (previous known as timestamps)
#
# Methods for calculating the ages (timestamps) of nodes in the tree given
# the branch lengths.

def get_tree_ages(tree, root=None, leaves=None, times=None, esp=0.001):
    """
    Use the branch lengths of a tree to set timestamps for each node
    Assumes ultrametric tree.

    Leaves have time 0
    """

    if root is None:
        root = tree.root

    if times is None:
        times = {}

    def walk(node):
        """helper function"""
        if node.is_leaf() or (leaves and node in leaves):
            t = times.get(node, 0.0)
        else:
            t2 = None
            for child in node.children:
                t = walk(child)

                # ensure branch lengths are ultrametrix
                if (t2 is not None) and (esp is not None):
                    assert abs(t - t2)/t < esp, (node.name, t, t2)
                t2 = t

        times[node] = t
        return t + node.dist
    walk(root)

    return times
get_tree_timestamps = get_tree_ages  # backwards compatiability


def set_dists_from_ages(tree, times):
    """
    Sets the branch lengths of a tree using a timestamp dict
    """

    for node in tree:
        if node.parent:
            node.dist = times[node.parent] - times[node]
        else:
            node.dist = 0.0
set_dists_from_timestamps = set_dists_from_ages  # backwards compatiability


def check_ages(tree, times):
    """Asserts that timestamps are consistent with tree"""

    for node in tree:
        if node.parent:
            if (times[node.parent] - times[node] < 0.0 \
                or abs(((times[node.parent] - times[node]) \
                         - node.dist)/node.dist) > .001):
                draw_tree_names(tree, maxlen=7, minlen=7)
                util.print_cols([(a.name, b) for a, b in times.items()])
                print
                print node.name, node.dist, times[node.parent] - times[node]
                raise Exception("negative time span")
check_timestamps = check_ages  # backwards compatiability


#=============================================================================
# parent tables

def tree2parent_table(tree, data_cols=None):
    """Converts tree to a parent table

    This parent table will have a special numbering for the internal nodes,
    such that their id is also their row in the table.

    parent table is a standard format of the Compbio Lab as of 02/01/2007.
    It is a list of triples (node_name, parent_name, dist, ...)

    * parent_name indicates the parent of the node.  If the node is a root
      (has no parent), then parent_name is -1

    * dist is the distance between the node and its parent.

    * additional columns can be added using the data_cols argument.  The
      values are looked up from node.data[col]
    """

    if data_cols is None:
        data_cols = []

    ptable = []

    for node in tree:
        if node.parent:
            pname = node.parent.name
        else:
            pname = -1
        row = [node.name, pname, node.dist]
        for col in data_cols:
            row.append(node.data[col])
        ptable.append(row)

    return ptable


def parent_table2tree(ptable, data_cols=None, convert_names=True):
    """Converts a parent table to a Tree

    if convert_names is True, names that are strings that look like integers
    are converted to ints.

    See tree2parent_table for details
    """

    if data_cols is None:
        data_cols = []

    tree = Tree()
    parents = {}

    # create nodes
    for row in ptable:
        name, parent = row[:2]
        if convert_names and name.isdigit():
            name = int(name)
        if convert_names and parent.isdigit() or parent == "-1":
            parent = int(parent)

        node = TreeNode(name)
        node.dist = row[2]
        tree.add(node)
        parents[node] = parent

        for col, val in zip(data_cols, row[3:]):
            node.data[col] = val

    # link up parents
    for node, parent_name in parents.iteritems():
        if parent_name == -1:
            tree.root = node
        else:
            parent = tree.nodes[parent_name]
            tree.add_child(parent, node)

    return tree


def tree2parent_table_ordered(tree, leaf_names=None):
    """Converts tree to a parent table

    This parent table will have a special numbering for the internal nodes,
    such that their id is also their row in the table.

    parent table is a standard format of the Compbio Lab as of 02/01/2007.
    It is a list of triples (node_name, parent_name, dist)

    * If the node is a leaf node_name is the leaf name (a string)
    * If the node is internal node_name is an int representing which row
      (0-based) the node is in the table.

    * parent_name indicates the parent of the node.  If the parent is root, a
      -1 is used as the parent_name.

    * dist is the distance between the node and its parent.

    Arguments:
    leaf_names -- specifies that a tree with only a subset of the leaves
                  should be used

    NOTE: root is not given a row, because root does not have a distance
    the nodeid of the root is -1
    """

    if leaf_names is not None:
        tree = subtree_by_leaf_names(tree, leaf_names, new_copy=True)
    else:
        leaf_names = tree.leaf_names()

    # assign a numbering to the leaves as specified
    nodeid = 0
    nodeids = {}
    nodes = []
    for leaf in leaf_names:
        nodeids[tree.nodes[leaf]] = leaf
        nodes.append(tree.nodes[leaf])
        nodeid += 1

    # assign a numbering to the internal nodes
    for node in tree:
        if node.is_leaf():
            continue
        if node == tree.root:
            nodeids[node] = -1
        else:
            nodeids[node] = nodeid
            nodeid += 1
            nodes.append(node)

    # make parent table
    parent_table = []
    for node in nodes:
        parent_table.append([nodeids[node], nodeids[node.parent], node.dist])

    return parent_table


def parent_table2tree_ordered(ptable):
    """Converts a parent table to a Tree

    See tree2parentTable for details
    """
    # TODO: allow named internal nodes

    tree = Tree()

    # create nodes
    maxint = 0
    for name, parent_name, dist in ptable:
        node = TreeNode(name)
        node.dist = dist
        tree.add(node)

        if isinstance(name, int):
            maxint = max(name, maxint)

    # make a root node
    tree.nextname = maxint + 1
    tree.make_root()

    # link up parents
    for name, parent_name, dist in ptable:
        if parent_name == -1:
            parent = tree.root
        else:
            parent = tree.nodes[parent_name]
        tree.add_child(parent, tree.nodes[name])

    return tree


def write_parent_table(ptable, out=sys.stdout):
    """Writes a parent table to out

    out can be a filename or file stream
    """
    out = util.open_stream(out, "w")
    for row in ptable:
        out.write("\t".join(map(str, row)) + "\n")


def read_parent_table(filename):
    """Reads a parent table from the file 'filename'

    filename can also be an open file stream
    """
    infile = util.open_stream(filename)
    ptable = []

    for line in infile:
        row = line.rstrip("\n").split("\t")
        name, parent, dist = row[:3]

        if name.is_digit():
            name = int(name)
        if parent.is_digit() or parent == "-1":
            parent = int(parent)

        ptable.append([name, parent, float(dist)] + row[3:])

    return ptable


#=============================================================================
# Tree visualization

def layout_tree(tree, xscale, yscale, minlen=-util.INF, maxlen=util.INF,
                rootx=0, rooty=0):
    """\
    Determines the x and y coordinates for every branch in the tree.

    Branch lengths are determined by node.dist
    """

    r"""
       /-----   ]
       |        ] nodept[node]
    ---+ node   ]
       |
       |
       \---------
    """

    # first determine sizes and nodepts
    coords = {}
    sizes = {}          # number of descendants (leaves have size 1)
    nodept = {}         # distance between node y-coord and top bracket y-coord

    def walk(node):
        """helper function"""
        # calculate new y-coordinate for node

        # compute node sizes
        sizes[node] = 0
        for child in node.children:
            sizes[node] += walk(child)

        if node.is_leaf():
            sizes[node] = 1
            nodept[node] = yscale - 1
        else:
            top = nodept[node.children[0]]
            bot = ((sizes[node] - sizes[node.children[-1]]) * yscale +
                   nodept[node.children[-1]])
            nodept[node] = (top + bot) / 2.0

        return sizes[node]
    walk(tree.root)

    # determine x, y coordinates
    def walk2(node, x, y):
        """helper function"""
        xchildren = x+min(max(node.dist*xscale, minlen), maxlen)
        coords[node] = [xchildren, y + nodept[node]]

        if not node.is_leaf():
            ychild = y
            for child in node.children:
                walk2(child, xchildren, ychild)
                ychild += sizes[child] * yscale
    walk2(tree.root, rootx, rooty)

    return coords


def layout_tree_hierarchical(tree, xscale, yscale,
                             rootx=0, rooty=0):
    """\
    Determines the x and y coordinates for every branch in the tree.

    Leaves are drawn to line up.  Best used for hierarchical clustering.
    """

    r"""
       /-----   ]
       |        ] nodept[node]
    ---+ node   ]
       |
       |
       \---------
    """

    # first determine sizes and nodepts
    coords = {}
    sizes = {}          # number of descendants (leaves have size 1)
    depth = {}          # how deep in tree is node
    nodept = {}         # distance between node y-coord and top bracket y-coord

    def walk(node):
        """helper function"""
        # calculate new y-coordinate for node

        # recurse: compute node sizes
        sizes[node] = 0
        for child in node.children:
            sizes[node] += walk(child)

        if node.is_leaf():
            sizes[node] = 1
            nodept[node] = yscale - 1
            depth[node] = 0
        else:
            top = nodept[node.children[0]]
            bot = ((sizes[node] - sizes[node.children[-1]]) * yscale +
                   nodept[node.children[-1]])
            nodept[node] = (top + bot) / 2.0
            depth[node] = max(depth[child] + 1 for child in node.children)

        return sizes[node]
    walk(tree.root)

    # determine x, y coordinates
    maxdepth = depth[tree.root]

    def walk2(node, x, y):
        """helper function"""
        xchildren = x + xscale * (maxdepth - depth[node])
        coords[node] = [xchildren, y + nodept[node]]

        if not node.is_leaf():
            ychild = y
            for child in node.children:
                walk2(child, x, ychild)
                ychild += sizes[child] * yscale
    walk2(tree.root, rootx, rooty)

    return coords


def layout_tree_vertical(layout, offset=None, root=0, leaves=None,
                         ydir=-1):
    """Make layout vertical"""

    if offset is None:
        if leaves is not None:
            for node in layout:
                if node.is_leaf():
                    offset = leaves - ydir*layout[node][0]
                    break
        else:
            for node in layout:
                if node.parent is None:
                    offset = root - ydir*layout[node][0]
                    break

    for node, (x, y) in layout.iteritems():
        layout[node] = [y, offset + ydir*x]
    return layout


#=============================================================================
# Tree color map

def tree_color_map(leafmap=lambda x: (0, 0, 0)):
    """Returns a simple color mixing colormap"""

    def null_func(tree):
        """helper function"""
        pass

    if leafmap is None:
        return null_func

    def func(tree):
        """helper function"""
        def walk(node):
            """helper function"""
            if node.is_leaf():
                node.color = leafmap(node)
            else:
                colors = []
                for child in node.children:
                    walk(child)
                    colors.append(child.color)
                node.color = color_mix(colors)
        walk(tree.root)
    return func


def color_mix(colors):
    """Mixes together several color vectors into one"""

    sumcolor = [0, 0, 0]
    for c in colors:
        sumcolor[0] += c[0]
        sumcolor[1] += c[1]
        sumcolor[2] += c[2]
    for i in range(3):
        sumcolor[i] /= float(len(colors))
    return sumcolor


def make_expr_mapping(maps, default_color=(0, 0, 0)):
    """Returns a function that maps strings matching an expression to a value

    maps -- a list of pairs (expr, value)
    """

    # find exact matches and expressions
    exacts = {}
    exps = []
    for key, val in maps:
        if "*" not in key:
            exacts[key] = val
        else:
            exps.append((key, val))

    # create mapping function
    def mapping(key):
        """helper function"""
        if key in exacts:
            return exacts[key]

        # return default color
        if not isinstance(key, str):
            return default_color

        # eval expressions first in order of appearance
        for exp, val in exps:
            if exp[-1] == "*":
                if key.startswith(exp[:-1]):
                    return val
            elif exp[0] == "*":
                if key.endswith(exp[1:]):
                    return val

        raise Exception("Cannot map key '%s' to any value" % key)
    return mapping


def read_tree_color_map(filename):
    """Reads a tree colormap from a file"""

    infile = util.open_stream(filename)
    maps = []

    for line in infile:
        expr, red, green, blue = line.rstrip().split("\t")
        maps.append([expr, map(float, (red, green, blue))])

    name2color = make_expr_mapping(maps)

    def leafmap(node):
        """helper function"""
        return name2color(node.name)

    return tree_color_map(leafmap)


#=========================================================================
# Draw Tree ASCII art

def draw_tree(tree, labels=None, scale=40, spacing=2, out=sys.stdout,
              canvas=None, x=0, y=0, display=True, label_offset=-1,
              minlen=1, maxlen=10000):
    """Print a ASCII Art representation of the tree"""

    if labels is None:
        labels = {}
    if canvas is None:
        canvas = textdraw.TextCanvas()

    xscale = scale
    yscale = spacing

    # determine node sizes
    sizes = {}
    nodept = {}

    def walk(node):
        """helper function"""
        if node.is_leaf():
            sizes[node] = 1
            nodept[node] = yscale - 1
        else:
            sizes[node] = 0
        for child in node.children:
            sizes[node] += walk(child)
        if not node.is_leaf():
            top = nodept[node.children[0]]
            bot = ((sizes[node] - sizes[node.children[-1]]) * yscale +
                   nodept[node.children[-1]])
            nodept[node] = (top + bot) / 2
        return sizes[node]
    walk(tree.root)

    def walk2(node, x, y):
        """helper function"""
        # calc coords
        xchildren = int(x + min(max(node.dist * xscale, minlen), maxlen))

        # draw branch
        canvas.line(x, y + nodept[node], xchildren, y + nodept[node], '-')
        if node.name in labels:
            branchlen = xchildren - x
            lines = str(labels[node.name]).split("\n")
            labelwidth = max(map(len, lines))

            labellen = min(labelwidth,
                           max(int(branchlen - 1), 0))
            canvas.text(x + 1 + (branchlen - labellen)/2.,
                        y + nodept[node] + label_offset,
                        labels[node.name], width=labellen)

        if node.is_leaf():
            canvas.text(xchildren + 1, y + yscale - 1, str(node.name))
        else:
            top = y + nodept[node.children[0]]
            bot = (y + (sizes[node] - sizes[node.children[-1]]) * yscale + nodept[node.children[-1]])

            # draw children
            canvas.line(xchildren, top, xchildren, bot, '|')

            ychild = y
            for child in node.children:
                walk2(child, xchildren, ychild)
                ychild += sizes[child] * yscale

            canvas.set(xchildren, y + nodept[node], '+')
            canvas.set(xchildren, top, '/')
            canvas.set(xchildren, bot, '\\')
        canvas.set(x, y + nodept[node], '+')
    walk2(tree.root, x + 0, 0)

    if display:
        canvas.display(out)


def draw_tree_lens(tree, *args, **kargs):
    """Draw tree with branch lengths"""
    labels = {}
    for node in tree.nodes.values():
        labels[node.name] = "%f" % node.dist

    draw_tree(tree, labels, *args, **kargs)


def draw_tree_boot_lens(tree, *args, **kargs):
    """Draw tree with bootstraps and branch lengths"""
    if not tree.has_data("boot"):
        draw_tree_lens(tree, *args, **kargs)
        return

    labels = {}
    for node in tree.nodes.values():
        if node.is_leaf():
            labels[node.name] = "%f" % node.dist
        else:
            if isinstance(node.data["boot"], int):
                labels[node.name] = "(%d) %f" % (node.data["boot"], node.dist)
            else:
                labels[node.name] = "(%.2f) %f" % (
                    node.data["boot"], node.dist)

    draw_tree(tree, labels, *args, **kargs)


def draw_tree_names(tree, *args, **kargs):
    """Draw tree with names"""
    labels = {}
    for node in tree.nodes.values():
        if not node.is_leaf():
            labels[node.name] = "%s" % node.name

    draw_tree(tree, labels, *args, **kargs)


def draw_tree_name_lens(tree, *args, **kargs):
    """Draw tree with names and branch lengths"""
    labels = {}
    for node in tree.nodes.values():
        if not node.is_leaf():
            labels[node.name] = "%s " % node.name
        else:
            labels[node.name] = ""
        labels[node.name] += "%f" % node.dist

    draw_tree(tree, labels, *args, **kargs)
