Cz staging

"""
ivy - a phylogenetics library and visual shell
http://www.reelab.net/ivy

Copyright 2010 Richard Ree <rree@fieldmuseum.org>

Required: ipython, matplotlib, scipy, numpy
Useful: dendropy, biopython, etc.
"""
## This program is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 3 of the
## License, or (at your option) any later version.

## This program is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.

## You should have received a copy of the GNU General Public License
## along with this program. If not, see
## <http://www.gnu.org/licenses/>.

import tree, layout, contrasts, ages
import bipart, genbank, nexus, newick, storage
#import nodearray, data
import treebase
#import db
#import contrib
try:
    import ltt as _ltt
    ltt = _ltt.ltt
except ImportError:
    pass

import chars, align, sequtil
## try: import vis
## except RuntimeError: pass

"""
Calculate node ages from branch lengths.

The function of interest is `ages2lengths`
"""

def ages2lengths(node, node_ages, results={}):
    """
    Convert node ages to branch lengths

    Args:
        node (Node): Node object
        node_ages (dict): Dict mapping nodes to ages
    Returns:
        dict: mapping of nodes to lengths

    """
    for d in node.descendants():
        age = node_ages[d]
        if d.parent:
            parent_age = node_ages[d.parent]
            results[d] = parent_age - age
    return results

def min_ages(node, leaf_ages, results={}):
    """
    Calculate minimum ages given fixed ages in leaf_ages

    Args:
        node (Node): A node object
        leaf_ages (dict): A dict mapping leaf nodes to ages
    Returns:
        dict: mapping of nodes to ages
    """
    v = []
    for child in node.children:
        if child.label and (child.label in leaf_ages):
            age = leaf_ages[child.label]
            v.append(age)
            results[child] = age
        else:
            min_ages(child, leaf_ages, results)
            age = results[child]
            v.append(age)
    results[node] = max(v)
    return results

def smooth(node, node_ages, results={}):
    """
    adjust ages of internal nodes by smoothing
    RR: I don't actually know what this function does -CZ
    """
    if node.parent:
        parent_age = node_ages[node.parent]
        if node.children:
            max_child_age = max([ node_ages[child] for child in node.children ])
            # make the new age the average of parent and max child
            new_node_age = (parent_age + max_child_age)/2.0
            results[node] = new_node_age
        else:
            results[node] = node_ages[node]
    else:
        results[node] = node_ages[node]
    for child in node.children:
        smooth(child, node_ages, results)
    return results

if __name__ == "__main__":
    import newick, ascii

    s = "((((a,b),(c,d),(e,f)),g),h);"
    root = newick.parse(s)

    leaf_ages = {
        "a": 3,
        "b": 2,
        "c": 4,
        "d": 1,
        "e": 3,
        "f": 0.5,
        "g": 10,
        "h": 5,
        }

    ma = min_ages(root, leaf_ages)
    d = ma
    for i in range(10):
        d = smooth(root, d)
    for node, val in ages2lengths(root, d).items():
        node.length = val
    print ascii.render(root, scaled=1)

import os
from subprocess import Popen, PIPE
from Bio import AlignIO
from Bio.Alphabet import IUPAC
from cStringIO import StringIO
from tempfile import NamedTemporaryFile

MUSCLE = "/usr/bin/muscle"

def muscle(seqs, cmd=None):
    if not cmd: cmd = MUSCLE
    assert os.path.exists(cmd)
    p = Popen([cmd], stdin=PIPE, stdout=PIPE)
    write = p.stdin.write
    for x in seqs:
        write(">%s\n%s\n" % (x.id, x.seq))
    out = p.communicate()[0]
    aln = AlignIO.read(StringIO(out), 'fasta', alphabet=IUPAC.ambiguous_dna)
    return aln

def musclep(seqs1, seqs2, cmd="/usr/bin/muscle"):
    assert os.path.exists(cmd)
    f1 = NamedTemporaryFile(); f2 = NamedTemporaryFile()
    for s, f in ((seqs1, f1), (seqs2, f2)):
        write = f.file.write
        for x in s: write(">%s\n%s\n" % (x.id, x.seq))
    f1.file.flush(); f2.file.flush()
    cmd += " -profile -in1 %s -in2 %s" % (f1.name, f2.name)
    p = Popen(cmd.split(), stdout=PIPE)
    out = p.communicate()[0]
    aln = AlignIO.read(StringIO(out), 'fasta', alphabet=IUPAC.ambiguous_dna)
    f1.file.close(); f2.file.close()
    return aln

def read(data, format=None, name=None):
    from types import StringTypes

    def strip(s):
        fname = os.path.split(s)[-1]
        head, tail = os.path.splitext(fname)
        tail = tail.lower()
        if tail in (".fasta", ".nex", ".nexus"):
            return head
        else:
            return fname

    if (not format):
        if (type(data) in StringTypes) and os.path.isfile(data):
            s = data.lower()
            if s.endswith("fasta"):
                format="fasta"
            for tail in ".nex", ".nexus":
                if s.endswith(tail):
                    format="nexus"
                    break

    if (not format):
        format = "fasta"

    if type(data) in StringTypes:
        if os.path.isfile(data):
            name = strip(data)
            with open(data) as f:
                return AlignIO.read(f, format, alphabet=IUPAC.ambiguous_dna)
        else:
            f = StringIO(data)
            return AlignIO.read(f, format, alphabet=IUPAC.ambiguous_dna)

    elif (hasattr(data, "tell") and hasattr(data, "read")):
        treename = strip(getattr(data, "name", None))
        return AlignIO.read(data, format, alphabet=IUPAC.ambiguous_dna)

    raise IOError, "unable to read alignment from '%s'" % data

def write(data, f, format='fasta'):
    AlignIO.write(data, f, format)

def find(aln, substr):
    """
    generator that yields (seqnum, pos) tuples for every position of
    ``subseq`` in `aln`
    """
    from sequtil import finditer
    N = len(substr)
    for i, rec in enumerate(aln):
        for j in finditer(rec.seq, substr):
            yield (i,j)

def find_id(aln, regexp):
    import re
    return [ (i,s) for i, s in enumerate(aln) if re.search(regexp, s.id) ]

def gapcols(aln, c='-'):
    from numpy import array
    a = array([ list(x.seq) for x in aln ])
    for i, col in enumerate(a.T):
        s = set(col==c)
        if len(s)==1 and True in s:
            yield i

from array import array
from layout import depth_length_preorder_traversal

class AsciiBuffer:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self._b = [ array('c', ' '*width) for line in range(height) ]

    def putstr(self, r, c, s):
        assert r < self.height
        assert c+len(s) <= self.width, "%s %s %s '%s'" % (self.width, r, c, s)
        self._b[r][c:c+len(s)] = array('c', s)

    def __str__(self):
        return "\n".join([ b.tostring() for b in self._b ])

def sum_to_root(node, internodes=True, length=False):
    """
    Number of branches from node to root.

    Args:
        node (Node): A Node object
        RR: Do internodes and length do anything in this function? -CZ
    Returns:
        int: The number of branches from node to root.
    """
    i = 0
    n = node
    while 1:
        if not n.parent:
            break
        else:
            n = n.parent
            i += 1
    return i

## def depth_length_preorder_traversal(node):
##     if not node.parent:
##         node.depth = 0
##         node.length_to_root = 0.0
##     else:
##         p = node.parent
##         node.depth = p.depth + 1
##         node.length_to_root = p.length_to_root + (node.length or 0.0)
##     for ch in node.children:
##         depth_length_preorder_traversal(ch)

def smooth_cpos(node, n2c):
    for ch in node.children:
        smooth_cpos(ch, n2c)

    if node.parent and not node.isleaf:
        px = n2c[node.parent].c
        cx = min([ n2c[ch].c for ch in node.children ])
        dxp = n2c[node].c - px
        cxp = cx - n2c[node].c
        node.c = int(px + (cx - px)*0.5)

def scale_cpos(node, n2c, scalef, root_offset):
    if node.parent:
        n2c[node].c = n2c[node.parent].c + int(node.length * scalef)
    else:
        n2c[node].c = root_offset

    for ch in node.children:
        scale_cpos(ch, n2c, scalef, root_offset)

def set_rpos(node, n2c):
    for child in node.children:
        set_rpos(child, n2c)
        nc = n2c[node]
        if node.children:
            children = node.children
            c0 = n2c[children[0]]
            c1 = n2c[children[-1]]
            rmin = c0.r; rmax = c1.r
            nc.r = int(rmin + (rmax-rmin)/2.0)

def render(root, unitlen=3, minwidth=50, maxwidth=None, scaled=False,
           show_internal_labels=True):
    """
    Create the ascii tree to be shown with print()
    """
    n2c = depth_length_preorder_traversal(root)
    leaves = root.leaves(); nleaves = len(leaves)
    maxdepth = max([ n2c[lf].depth for lf in leaves ])
    max_labelwidth = max([ len(lf.label) for lf in leaves ]) + 1

    root_offset = 0
    if root.label and show_internal_labels:
        root_offset = len(root.label)

    width = maxdepth*unitlen + max_labelwidth + 2 + root_offset
    height = 2*nleaves - 1

    if width < minwidth:
        unitlen = (minwidth - max_labelwidth - 2 - root_offset)/maxdepth
        width = maxdepth*unitlen + max_labelwidth + 2 + root_offset

    buf = AsciiBuffer(width, height)

    for i, lf in enumerate(leaves):
        c = n2c[lf]
        c.c = width - max_labelwidth - 2
        c.r = i*2

    for node in root.postiter():
        nc = n2c[node]
        if node.children:
            children = node.children
            c0 = n2c[children[0]]
            c1 = n2c[children[-1]]
            rmin = c0.r; rmax = c1.r
            nc.r = int(rmin + (rmax-rmin)/2.0)
            nc.c = min([ n2c[ch].c for ch in children ]) - unitlen

    if not scaled:
        smooth_cpos(root, n2c)
    else:
        maxlen = max([ n2c[lf].length_to_root for lf in leaves ])
        scalef = (n2c[leaves[0]].c + 1 - root_offset)/maxlen
        scale_cpos(root, n2c, scalef, root_offset)

    for node in root.postiter():
        nc = n2c[node]
        if node.parent:
            pc = n2c[node.parent]
            for r in range(min([nc.r, pc.r]),
                           max([nc.r, pc.r])):
                buf.putstr(r, pc.c, ":")

            sym = getattr(nc, "hchar", "-")
            vbar = sym*(nc.c-pc.c)
            buf.putstr(nc.r, pc.c, vbar)

        if node.isleaf:
            buf.putstr(nc.r, nc.c+1, " "+node.label)
        else:
            if node.label and show_internal_labels:
                buf.putstr(nc.r, nc.c-len(node.label), node.label)

        buf.putstr(nc.r, nc.c, "+")

    return str(buf)

if __name__ == "__main__":
    import random, tree
    rand = random.Random()

    t = tree.read(
        "(foo,((bar,(dog,cat)dc)dcb,(shoe,(fly,(cow, bowwow)cowb)cbf)X)Y)Z;"
        )

    #t = tree.read("(((foo:4.6):5.6, (bar:6.5, baz:2.3):3.0):3.0);")
    #t = tree.read("(foo:4.6, (bar:6.5, baz:2.3)X:3.0)Y:3.0;")

    i = 1
    print render(t, scaled=0, show_internal_labels=1)
    r = t.get("cat").parent
    tree.reroot(t, r)
    tp = t.parent
    tp.remove_child(t)
    c = t.children[0]
    t.remove_child(c)
    tp.add_child(c)
    print render(r, scaled=0, show_internal_labels=1)

"""
For drawing big trees.  Calculate which clades can be 'collapsed' and
displayed with a placeholder.

TODO: test and develop this module further
"""
from storage import Storage

def autocollapse_info(node, collapsed, visible=True, info={}):
    """
    gather information to determine if a node should be collapsed

    *collapsed* is a set containing nodes that are already collapsed
    """
    if node not in info:
        s = Storage()
        info[node] = s
    else:
        s = info[node]

    if visible and (node in collapsed):
        visible = False

    nnodes = 1 # total number of nodes, including node
    # number of visible leaves
    nvisible = int((visible and node.isleaf) or (node in collapsed))
    ntips = int(node.isleaf)
    ntips_visible = int(node.isleaf and visible)
    s.has_labeled_descendant = False
    s.depth = 1

    for child in node.children:
        autocollapse_info(child, collapsed, visible, info)
        cs = info[child]
        nnodes += cs.nnodes
        nvisible += cs.nvisible
        ntips += cs.ntips
        ntips_visible += cs.ntips_visible
        if (child.label and (not child.isleaf)) \
           or (cs.has_labeled_descendant):
            s.has_labeled_descendant = True
        if cs.depth >= s.depth:
            s.depth = cs.depth+1
    s.nnodes = nnodes
    s.nvisible = nvisible
    s.ntips = ntips
    s.ntips_visible = ntips_visible
    return info

def autocollapse(root, collapsed=None, keep_visible=None, max_visible=1000):
    """
    traverse a tree and find nodes that should be collapsed in order
    to satify *max_visible*

    *collapsed* is a set object for storing collapsed nodes

    *keep_visible* is a set object of nodes that should not be placed
    in *collapsed*
    """
    collapsed = collapsed or set()
    keep_visible = keep_visible or set()
    ntries = 0
    while True:
        if ntries > 10:
            return
        info = autocollapse_info(root, collapsed)
        nvisible = info[root].nvisible
        if nvisible <= max_visible:
            return

        v = []
        for node in root.iternodes():
            s = info[node]
            if (node.label and (not node.isleaf) and node.parent and
                (node not in keep_visible)):
                w = s.nvisible/float(s.depth)
                if s.has_labeled_descendant:
                    w *= 0.25
                v.append((w, node, s))
        v.sort(); v.reverse()
        for w, node, s in v:
            if node not in keep_visible and s.nvisible < (nvisible-1):
                print node
                collapsed.add(node)
                nvisible -= s.nvisible
            if nvisible <= max_visible:
                break
        ntries += 1
    return collapsed

import sys
from pprint import pprint
from glob import glob
from storage import Storage
from collections import defaultdict

## class BipartSet(object):
##     "A set of bipartitions"
##     def __init__(self, elements):
##         self.elements = frozenset(elements)
##         self.ref = sorted(elements)[0]
##         self.node2bipart = Storage()

##     def add(self, subset, node):
##         # filter out elements of subset not in 'elements'
##         subset = (frozenset(subset) & self.elements)
##         if self.ref not in self.subset:
##             self.subset = self.elements - self.subset

class Bipart(object):
    """
    A class representing a bipartition.
    """
    def __init__(self, elements, subset, node=None, support=None):
        """
        'elements' and 'subset' are set objects
        """
        self.subset = subset
        self.compute(elements)
        self.node = node
        self.support = support

    def __hash__(self):
        return self._hash

    def __eq__(self, other):
        assert self.elements == other.elements
        return ((self.subset == other.subset) or
                (self.subset == (self.elements - other.subset)))

    def __repr__(self):
        v = sorted(self.subset)
        return "(%s)" % " ".join(map(str, v))

    def compute(self, elements):
        self.elements = frozenset(elements)
        self.ref = sorted(elements)[0]
        # filter out elements of subset not in 'elements'
        self.subset = (frozenset(self.subset) & self.elements)
        self._hash = hash(self.subset)
        if self.ref not in self.subset:
            self.subset = self.elements - self.subset
        self.complement = self.elements - self.subset

    def iscompatible(self, other):
        ## assert self.elements == other.elements
        if (self.subset.issubset(other.subset) or
            other.subset.issubset(self.subset)):
            return True
        if (((self.subset | other.subset) == self.elements) or
            (not (self.subset & other.subset))):
            return True
        return False

def conflict(bp1, bp2, support=None):
    if ((support and (bp1.support >= support) and (bp2.support >= support))
        or (not support)):
        if not bp1.iscompatible(bp2):
            return True
    return False

class TreeSet:
    def __init__(self, root, elements=None):
        self.root = root
        self.node2labels = root.leafsets(labels=True)
        self.elements = elements or self.node2labels.pop(root)
        self.biparts = [ Bipart(self.elements, v, node=k,
                                support=int(k.label or 0))
                         for k, v in self.node2labels.items() ]

def compare_trees(r1, r2, support=None):
    e = (set([ x.label for x in r1.leaves() ]) &
         set([ x.label for x in r2.leaves() ]))
    bp1 = [ Bipart(e, v, node=k, support=int(k.label or 0))
            for k, v in r1.leafsets(labels=True).items() ]
    bp2 = [ Bipart(e, v, node=k, support=int(k.label or 0))
            for k, v in r2.leafsets(labels=True).items() ]
    return compare(bp1, bp2, support)

def compare(set1, set2, support=None):
    hits1 = []; hits2 = []
    conflicts1 = defaultdict(set); conflicts2 = defaultdict(set)
    for bp1 in set1:
        for bp2 in set2:
            if bp1 == bp2:
                hits1.append(bp1.node); hits2.append(bp2.node)
            if conflict(bp1, bp2, support):
                conflicts1[bp1.node].add(bp2.node)
                conflicts2[bp2.node].add(bp1.node)
    return hits1, hits2, conflicts1, conflicts2

## a = Bipart("abcdef", "abc")
## b = Bipart("abcdef", "def")
## c = Bipart("abcdef", "ab")
## d = Bipart("abcdef", "cd")
## print a == b
## print a.iscompatible(b)
## print a.iscompatible(c)
## print a.iscompatible(d)
## print c.iscompatible(d)
## sys.exit()

import scipy, numpy

def default_costmatrix(numstates, dtype=numpy.int):
    "a square array with zeroes along the diagonal, ones elsewhere"
    return scipy.logical_not(scipy.identity(numstates)).astype(float)

def minstates(v):
    "return the indices of v that equal the minimum"
    return scipy.nonzero(scipy.equal(v, min(v)))

def downpass(node, states, stepmatrix, chardata, node2dpv=None):
    if node2dpv is None:
        node2dpv = {}

    if not node.isleaf:
        for child in node.children:
            downpass(child, states, stepmatrix, chardata, node2dpv)

        dpv = scipy.zeros([len(states)])
        node2dpv[node] = dpv
        for i in states:
            for child in node.children:
                child_dpv = node2dpv[child]
                mincost = min([ child_dpv[j] + stepmatrix[i,j] \
                                for j in states ])
                dpv[i] += mincost

        #print node.label, node.dpv

    else:
        #print node.label, chardata[node.label]
        node2dpv[node] = stepmatrix[:,chardata[node.label]]

    return node2dpv

def uppass(node, states, stepmatrix, node2dpv, node2upm={},
           node2ancstates=None):
    parent = node.parent
    if not node.isleaf:
        if parent is None: # root
            dpv = node2dpv[node]
            upm = None
            node.mincost = min(dpv)
            node2ancstates = {node: minstates(dpv)}

        else:
            M = scipy.zeros(stepmatrix.shape)
            for i in states:
                sibs = [ c for c in parent.children if c is not node ]
                for j in states:
                    c = 0
                    for sib in sibs:
                        sibdpv = node2dpv[sib]
                        c += min([ sibdpv[x] + stepmatrix[j,x]
                                   for x in states ])
                    c += stepmatrix[j,i]

                    p_upm = node2upm.get(parent)
                    if p_upm is not None:
                        c += min(p_upm[j])

                    M[i,j] += c

            node2upm[node] = M

            v = node2dpv[node][:]
            for s in states:
                v[s] += min(M[s])
            node2ancstates[node] = minstates(v)

        for child in node.children:
            uppass(child, states, stepmatrix, node2dpv, node2upm,
                   node2ancstates)

    return node2ancstates

def ancstates(tree, chardata, stepmatrix):
    states = range(len(stepmatrix))
    return uppass(tree, states, stepmatrix,
                  downpass(tree, states, stepmatrix, chardata))

def _bindeltran(node, stepmatrix, node2dpv, node2deltr=None, ancstate=None):
    if node2deltr is None:
        node2deltr = {}

    dpv = node2dpv[node]
    if ancstate is not None:
        c, s = min([ (cost+stepmatrix[ancstate,i], i) \
                     for i, cost in enumerate(dpv) ])
    else:
        c, s = min([ (cost, i) for i, cost in enumerate(dpv) ])

    node2deltr[node] = s
    for child in node.children:
        _bindeltran(child, stepmatrix, node2dpv, node2deltr, s)

    return node2deltr

def binary_deltran(tree, chardata, stepmatrix):
    states = range(len(stepmatrix))
    node2dpv = downpass(tree, states, stepmatrix, chardata)
    node2deltr = _bindeltran(tree, stepmatrix, node2dpv)
    return node2deltr

if __name__ == "__main__":
    from pprint import pprint
    from ivy import tree
    root = tree.read("(a,((b,c),(d,(e,f))));")

    nstates = 4
    states = range(nstates)
    cm = default_costmatrix(nstates)
    chardata = dict(zip("abcdef", map(int, "000233")))
    dp = downpass(root, states, cm, chardata)

    for i, node in enumerate(root):
        if not node.label:
            node.label = "N%s" % i
        else:
            node.label = "%s (%s)" % (node.label, chardata[node.label])

    print ascii.render(root)

##     nstates = 2
##     leaves = tree.leaves()
##     for leaf in leaves:
##         leaf.anc_cost_vector = chardata[leaf.label]

    pprint(
        #ancstates(root, chardata, cm)
        #uppass(root, states, cm, downpass(tree, states, cm, chardata))
        dp
        )

#!/usr/bin/env python
"""
Functions for evolving traits and trees.
"""
def brownian(root, sigma=1.0, init=0.0, values={}):
    """
    Recursively evolve a trait by Brownian motion up from the node
    *root*.

    * *sigma*: standard deviation of the normal random variate after
      one unit of branch length

    * *init*: initial value

    Returns: *values* - a dictionary mapping nodes to evolved values
    """
    from scipy.stats import norm
    values[root] = init
    for child in root.children:
        time = child.length
        random_step = norm.rvs(init, scale=sigma*time)
        brownian(child, sigma, random_step, values)
    return values

def test_brownian():
    """
    Evolve a trait up an example tree of primates:.

    ((((Homo:0.21,Pongo:0.21)N1:0.28,Macaca:0.49)N2:0.13,
    Ateles:0.62)N3:0.38,Galago:1.00)root;

    Returns: (*root*, *data*) - the root node and evolved data.
    """
    import newick
    root = newick.parse(
        "((((Homo:0.21,Pongo:0.21)N1:0.28,Macaca:0.49)N2:0.13,"\
        "Ateles:0.62)N3:0.38,Galago:1.00)root;"
        )
    print root.ascii(scaled=True)
    evolved = brownian(root)
    for node in root.iternodes():
        print node.label, evolved[node]
    return root, evolved

if __name__ == "__main__":
    test_brownian()

"""
Categorical Markov models with k states.
"""
import numpy, scipy, random
import scipy.linalg
import scipy.optimize
from scipy import array, zeros, ones
from scipy.linalg import expm#, expm2, expm3
from math import log, exp
rand = random.Random()
uniform = rand.uniform; expovariate = rand.expovariate

LARGE = 10e10 # large -lnL value used to bound parameter optimization

class Q:
    def __init__(self, k=2, layout=None):
        """
        Represents a square transition matrix with k states.

        'layout' is a square (k,k) array of integers that index free
        rate parameters (values on the diagonal are ignored).  Cells
        with value 0 will have the first rate parameter, 1 the
        second, etc.
        """
        self.k = k
        self.range = range(k)
        self.offdiag = array(numpy.eye(k)==0, dtype=numpy.int)
        if layout is None:
            layout = zeros((k,k), numpy.int)
        self.layout = layout*self.offdiag

    def fill(self, rates):
        m = numpy.take(rates, self.layout)*self.offdiag
        v = m.sum(1) * -1
        for i in self.range:
            m[i,i] = v[i]
        return m

    def default_priors(self):
        p = 1.0/self.k
        return [p]*self.k

def sample_weighted(weights):
    u = uniform(0, sum(weights))
    x = 0.0
    for i, w in enumerate(weights):
        x += w
        if u < x:
            break
    return i

def conditionals(root, data, Q):
    nstates = Q.shape[0]
    states = range(nstates)
    nodes = [ x for x in root.postiter() ]
    nnodes = len(nodes)
    v = zeros((nnodes,nstates))
    n2i = {}

    for i, n in enumerate(nodes):
        n2i[n] = i
        if n.isleaf:
            state = data[n.label]
            try:
                state = int(state)
                v[i,state] = 1.0
            except ValueError:
                if state == '?' or state == '-':
                    v[i,:] += 1/float(nstates)
        else:
            Pv = [ (expm(Q*child.length)*v[n2i[child]]).sum(1)
                   for child in n.children ]
            v[i] = numpy.multiply(*Pv)
            # fossils
            state = None
            if n.label in data:
                state = int(data[n.label])
            elif n in data:
                state = int(data[n])
            if state != None:
                for s in states:
                    if s != state: v[i,s] = 0.0

    return dict([ (n, v[i]) for n,i in n2i.items() ])

def contrasts(root, data, Q):
    cond = conditionals(root, data, Q)
    d = {}
    for n in root.postiter(lambda x:x.children):
        nc = cond[n]; nc /= sum(nc)
        diff = 0.0
        for child in n.children:
            cc = cond[child]; cc /= sum(cc)
            diff += numpy.sum(numpy.abs(nc-cc))
        d[n] = diff
    return d

def lnL(root, data, Q, priors):
    d = conditionals(root, data, Q)
    return numpy.log(sum(d[root]*priors))

def optimize(root, data, Q, priors=None):
    Qfill = Q.fill
    if priors is None: priors = Q.default_priors()
    def f(params):
        if (params<0).any(): return LARGE
        return -lnL(root, data, Qfill(params), priors)

    # initial parameter values
    p = [1.0]*len(set(Q.layout.flat))

    v = scipy.optimize.fmin_powell(
        f, p, full_output=True, disp=0, callback=None
        )
    params, neglnL = v[:2]
    if neglnL == LARGE:
        raise Exception("ConvergenceError")
    return params, neglnL

def sim(root, n2p, s0, d=None):
    if d is None:
        d = {root:s0}
    for n in root.children:
        v = n2p[n][s0]
        i = sample_weighted(v)
        d[n] = i
        sim(n, n2p, i, d)
    return d

def stmap(root, states, ancstates, Q, condition_on_success):
    """
    This and its dependent functions below need testing and
    optimization.
    """
    results = []
    for n in root.descendants():
        si = ancstates[n.parent]
        sj = ancstates[n]
        v = simulate_on_branch(states, si, sj, Q, n.length,
                               condition_on_success)
        print n, si, sj
        if v:
            results.append(v)
        else:
            return None
    return results

def simulate_on_branch(states, si, sj, Q, brlen, condition_on_success):
    point = 0.0
    history = [(si, point)]
    if si != sj:  # condition on one change occurring
        lambd = -(Q[si,si])
        U = uniform(0.0, 1.0)
        # see appendix of Nielsen 2001, Genetics
        t = brlen - point
        newpoint = -(1.0/lambd) * log(1.0 - U*(1.0 - exp(-lambd * t)))
        newstate = draw_new_state(states, Q, si)
        history.append((newstate, newpoint))
        si = newstate; point = newpoint
    while 1:
        lambd = -(Q[si,si])
        rv = expovariate(lambd)
        newpoint = point + rv

        if newpoint <= brlen:  # state change along branch
            newstate = draw_new_state(states, Q, si)
            history.append((newstate, newpoint))
            si = newstate; point = newpoint
        else:
            history.append((si, brlen))
            break

    if si == sj or (not condition_on_success): # success
        return history

    return None

def draw_new_state(states, Q, si):
    """
    Given a rate matrix Q, a starting state si, and an ordered
    sequence of states, eg (0, 1), draw a new state sj with
    probability -(qij/qii)
    """
    Qrow = Q[si]
    qii = Qrow[si]
    qij_probs = [ (x, -(Qrow[x]/qii)) for x in states if x != si ]
    uni = uniform(0.0, 1.0)
    val = 0.0
    for sj, prob in qij_probs:
        val += prob
        if uni < val:
            return sj

def sample_ancstates(node, states, conditionals, n2p, fixed={}):
    """
    Sample ancestral states from their conditional likelihoods
    """
    ancstates = {}
    for n in node.preiter():
        if n in fixed:
            state = fixed[n]
        else:
            cond = conditionals[n]

            if n.parent:
                P = n2p[n]
                ancst = ancstates[n.parent]
                newstate_Prow = P[ancst]
                cond *= newstate_Prow

            cond /= sum(cond)

            rv = uniform(0.0, 1.0)
            v = 0.0
            for state, c in zip(states, cond):
                v += c
                if rv < v:
                    break
        ancstates[n] = state

    return ancstates

"""
Calculate independent contrasts

TODO: include utilities for transforming data, etc.
"""
def PIC(node, data, results={}):
    """
    Phylogenetic independent contrasts.

    Recursively calculate independent contrasts of a bifurcating node
    given a dictionary of trait values.

    Args:
        node (Node): A node object
        data (dict): Mapping of leaf names to character values

    Returns:
        dict: Mapping of internal nodes to tuples containing ancestral
              state, its variance (error), the contrast, and the
              contrasts's variance.

    TODO: modify to accommodate polytomies.
    """
    X = []; v = []
    for child in node.children:
        if child.children:
            PIC(child, data, results)
            child_results = results[child]
            X.append(child_results[0])
            v.append(child_results[1])
        else:
            X.append(data[child.label])
            v.append(child.length)

    Xi, Xj = X  # Xi - Xj is the contrast value
    vi, vj = v

    # Xk is the reconstructed state at the node
    Xk = ((1.0/vi)*Xi + (1/vj)*Xj) / (1.0/vi + 1.0/vj)

    # vk is the variance
    vk = node.length + (vi*vj)/(vi+vj)

    results[node] = (Xk, vk, Xi-Xj, vi+vj)

    return results

if __name__ == "__main__":
    import tree
    n = tree.read(
        "((((Homo:0.21,Pongo:0.21)N1:0.28,Macaca:0.49)N2:0.13,"\
        "Ateles:0.62)N3:0.38,Galago:1.00)N4:0.0;"
        )
    char1 = {
        "Homo": 4.09434,
        "Pongo": 3.61092,
        "Macaca": 2.37024,
        "Ateles": 2.02815,
        "Galago": -1.46968
        }

    for k, v in PIC(n, char1).items():
        print k.label or k.id, v

import re, sys, logging
from collections import defaultdict
from itertools import izip_longest, ifilter
from Bio import Entrez, SeqIO
from Bio.Blast import NCBIWWW, NCBIXML
from ivy.storage import Storage

email = ""

def batch(iterable, size):
    """
    Take an iterable and return it in chunks (sub-iterables)

    Args:
        iterable: Any iterable
        size (int): Size of chunks
    Yields:
        Chunks of size `size`
    """
    args = [ iter(iterable) ]*size
    for x in izip_longest(fillvalue=None, *args):
        yield ifilter(None, x)

def extract_gbac(s):
    """
    Extract genbank accession

    Args:
        s (str): text string of genbank file
    Returns:
        list: Accession number(s)
    """
    gbac_re = re.compile(r'[A-Z]{1,2}[0-9]{4,7}')
    return gbac_re.findall(s, re.M)
    # RR: This also returns various other strings that match the pattern (eg.
    # protein ids)

def extract_gene(seq, gene):
    """
    RR: Not sure what format seq should be in -CZ
    """
    for t in "exon", "gene":
        for x in seq.features:
            if x.type == t:
                v = x.qualifiers.get("gene")
                if v == [gene]:
                    if x.sub_features:
                        s = [ seq[sf.location.start.position:
                                  sf.location.end.position]
                              for sf in x.sub_features ]
                        return reduce(lambda x,y:x+y, s)
                    else:
                        loc = x.location
                        return seq[loc.start.position-10:loc.end.position+10]

def gi2webenv(gilist):
    h = Entrez.esearch(
        db="nucleotide", term=" OR ".join(gilist), usehistory="y",
        retmax=len(gilist)
        )
    d = Entrez.read(h)
    return d["WebEnv"], d["QueryKey"]

def gi2tax(gi):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.elink(dbfrom='taxonomy', db='nucleotide', from_uid=gi,
                     LinkName='nucleotide_taxonomy')
    r = Entrez.read(h)[0]
    h.close()
    i = r['LinkSetDb'][0]['Link'][0]['Id']
    h = Entrez.efetch(db='taxonomy', id=i, retmode='xml')
    r = Entrez.read(h)[0]
    h.close()
    return r

def ac2gi(ac):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.esearch(db="nucleotide", term=ac, retmax=1)
    d = Entrez.read(h)['IdList'][0]
    h.close()
    return d

def acsum(aclist, batchsize=100):
    """
    fetch esummary info for list of accession numbers -- useful for
    getting gi and taxids
    """
    global email
    assert email, "set email!"
    Entrez.email = email
    results = {}
    for v in batch(aclist, batchsize):
        v = list(v)
        h = Entrez.esearch(
            db="nucleotide", retmax=len(v),
            term=" OR ".join([ "%s[ACCN]" % x for x in v ]),
            usehistory="y"
            )
        d = Entrez.read(h)
        h.close()
        # gis, but not in order of aclist
        gis = d['IdList']
        d = Entrez.read(Entrez.esummary(db='nucleotide', id=','.join(gis)),
                        validate=False)
        for x in d:
            ac = x['Caption']
            if ac in aclist:
                results[ac] = x
    return results

def fetch_aclist(aclist, batchsize=1000):
    global email
    assert email, "set email!"
    Entrez.email = email
    results = {}
    n = 0
    for v in batch(aclist, batchsize):
        v = list(v)
        h = Entrez.esearch(
            db="nucleotide",
            term=" OR ".join([ "%s[ACCN]" % x for x in v ]),
            usehistory="y"
            )
        d = Entrez.read(h)
        h.close()
        h = Entrez.efetch(db="nucleotide", rettype="gb", retmax=len(v),
                          webenv=d["WebEnv"], query_key=d["QueryKey"])
        seqs = SeqIO.parse(h, "genbank")
        for s in seqs:
            try:
                ac = s.annotations["accessions"][0]
                if ac in aclist:
                    results[ac] = s
            except:
                pass
        h.close()
        n += len(v)
        logging.info('fetched %s sequences', n)
    return results

def fetch_gilist(gilist, batchsize=1000):
    global email
    assert email, "set email!"
    Entrez.email = email
    results = {}
    for v in batch(gilist, batchsize):
        v = map(str, v)
        h = Entrez.epost(db="nucleotide", id=",".join(v), usehistory="y")
        d = Entrez.read(h)
        h.close()
        h = Entrez.efetch(db="nucleotide", rettype="gb", retmax=len(v),
                          webenv=d["WebEnv"], query_key=d["QueryKey"])
        seqs = SeqIO.parse(h, "genbank")
        for s in seqs:
            try:
                gi = s.annotations["gi"]
                if gi in v:
                    s.id = organism_id(s)
                    results[gi] = s
            except:
                pass
        h.close()
    return results

def organism_id(s):
    org = (s.annotations.get('organism') or '').replace('.', '')
    return '%s_%s' % (org.replace(' ','_'), s.id.split('.')[0])

def fetchseq(gi):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.efetch(db="nucleotide", id=str(gi), rettype="gb")
    s = SeqIO.read(h, 'genbank')
    s.id = organism_id(s)
    return s

def create_fastas(data, genes):
    fastas = dict([ (g, file(g+".fasta", "w")) for g in genes ])
    for label, seqs in data.items():
        for gene, s in zip(genes, seqs):
            if s and type(s) != str:
                tag = None
                try:
                    tag = "%s_%s" % (label, s.annotations["accessions"][0])
                except:
                    tag = "%s_%s" % (label, s.name)
                if tag:
                    fastas[gene].write(">%s\n%s\n" % (tag, s.seq))
            else:
                sys.stderr.write(("error: not an accession number? "
                                  "%s (%s %s)\n" % (s, label, gene)))

    for f in fastas.values(): f.close()

def merge_fastas(fnames, name="merged"):
    outfile = file(name+".phy", "w")
    gene2len = {}
    d = defaultdict(dict)
    for fn in fnames:
        gene = fn.split(".")[0]
        for rec in SeqIO.parse(file(fn), "fasta"):
            #sp = "_".join(rec.id.split("_")[:2])
            if rec.id.startswith("Pedicularis"):
                sp = rec.id.split("_")[1]
            else:
                sp = rec.id.split("_")[0]
            sp = "_".join(rec.id.split("_")[:-1])
            seq = str(rec.seq)
            d[sp][gene] = seq
            if gene not in gene2len:
                gene2len[gene] = len(seq)

    ntax = len(d)
    nchar = sum(gene2len.values())
    outfile.write("%s %s\n" % (ntax, nchar))
    genes = list(sorted(gene2len.keys()))
    for sp, data in sorted(d.items()):
        s = "".join([ (data.get(gene) or "".join(["?"]*gene2len[gene]))
                      for gene in genes ])
        outfile.write("%s  %s\n" % (sp, s))
    outfile.close()
    parts = file(name+".partitions", "w")
    i = 1
    for g in genes:
        n = gene2len[g]
        parts.write("DNA, %s = %s-%s\n" % (g, i, i+n-1))
        i += n
    parts.close()

def blast_closest(fasta, e=10):
    f = NCBIWWW.qblast("blastn", "nr", fasta, expect=e, hitlist_size=1)
    rec = NCBIXML.read(f)
    d = rec.descriptions[0]
    result = Storage()
    gi = re.findall(r'gi[|]([0-9]+)', d.title) or None
    if gi: result.gi = int(gi[0])
    ac = re.findall(r'gb[|]([^|]+)', d.title) or None
    if ac: result.ac = ac[0].split(".")[0]
    result.title = d.title.split("|")[-1].strip()
    return result

def blast(query, e=10, n=100, entrez_query=""):
    f = NCBIWWW.qblast("blastn", "nr", query, expect=e, hitlist_size=n,
                       entrez_query=entrez_query)
    recs = NCBIXML.parse(f)
    return recs
    ## v = []
    ## for d in rec.descriptions:
    ##     result = Storage()
    ##     gi = re.findall(r'gi[|]([0-9]+)', d.title) or None
    ##     if gi: result.gi = int(gi[0])
    ##     ac = re.findall(r'gb[|]([^|]+)', d.title) or None
    ##     if ac: result.ac = ac[0].split(".")[0]
    ##     result.title = d.title.split("|")[-1].strip()
    ##     v.append(result)
    ## return v

def start_codons(seq):
    i = seq.find('ATG')
    while i != -1:
        yield i
        i = seq.find('ATG', i+3)

def search_taxonomy(q):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.esearch(db="taxonomy", term=q)
    return Entrez.read(h)['IdList']

def fetchtax(taxid):
    global email
    assert email, "set email!"
    Entrez.email = email
    n = 1
    if not isinstance(taxid, int):
        # string, possibly with multiple values?
        try:
            taxid = taxid.strip()
            n = taxid.count(',') + 1
        except AttributeError:
            # iterable of values?
            try:
                n = len(taxid)
                taxid = ','.join(map(str, taxid))
            except TypeError:
                pass
    else:
        taxid = str(taxid)
    h = Entrez.efetch(db='taxonomy', id=taxid, retmode='xml', retmax=n)
    if n == 1:
        r = Entrez.read(h)[0]
    else:
        # a list of taxonomy results in same order of taxids
        r = Entrez.read(h)
    return r

__FIRST = re.compile('[^-]')
__LAST = re.compile('[-]*$')
def trimpos(rec):
    'return the positions of the first and last ungapped base'
    s = rec.seq.tostring()
    first = __FIRST.search(s).start()
    last = __LAST.search(s).start()-1
    return (first, last)

def fetch_DNA_seqs(terms, maxn=10000, batchsize=1000):
    """
    terms: sequence of search terms, quoted appropriately, with Entrez
           specifiers, e.g. ['"Mus musculus"[organism]']
    maxn: maximum number of sequences to return
    returns list of SeqRecord objects
    """
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.esearch(db="nucleotide", term=" OR ".join(terms), usehistory="y")
    d = Entrez.read(h)
    env = d['WebEnv']; key = d['QueryKey']
    N = int(d['Count'])
    if maxn: N = min(N, maxn)
    logging.info('fetching %s sequences', N)
    retstart = 0
    seqs = []
    n = 0
    while n < N:
        h = Entrez.efetch(
            db="nucleotide", rettype='gb', webenv=env, query_key=key,
            retstart=retstart, retmax=batchsize
            )
        v = list(SeqIO.parse(h, "genbank"))
        n += len(v)
        logging.info('...fetched %s', n)
        seqs.extend(v)
        retstart += batchsize
    logging.info('...done')
    return seqs

def seqrec_taxid(seqrec):
    "extract the NCBI taxon id from a sequence record"
    for ft in seqrec.features:
        if ft.type == 'source':
            break
    try:
        for x in ft.qualifiers['db_xref']:
            if x.startswith('taxon:'):
                return int(x.split(':')[1])
    except:
        pass

#!/usr/bin/env ipython
# -*- coding: utf-8 -*-
"""
Adds to the interactive IPython/pylab environment
"""
import sys, os, re
import ivy
import ivy.vis
from ivy.vis import symbols

def readtree(data, *args, **kwargs): return ivy.tree.read(data, *args, **kwargs)

def readaln(data, *args, **kwargs): return ivy.align.read(data, *args, **kwargs)

def treefig(*args, **kwargs):
    from ivy.vis import TreeFigure, MultiTreeFigure
    if len(args) == 1:
        fig = TreeFigure(args[0], **kwargs)
    else:
        fig = MultiTreeFigure(**kwargs)
        for arg in args:
            print arg
            fig.add(arg)
    fig.show()
    return fig

def alnfig(*args, **kwargs):
    from ivy.vis import AlignmentFigure
    return AlignmentFigure(*args, **kwargs)

def __maketree(self, s):
    import os#, IPython
    words = s.split()
    treename = "root"
    fname = None
    if words:
        treename = words.pop(0)
        if words and os.path.isfile(words[0]):
            fname = words.pop(0)

    if not fname:
        ## msg = "\n".join([
        ##     "Name of tree file",
        ##     "(Try dragging one into the terminal):\n"
        ##     ])
        msg = "Enter the name of a tree file or a newick string:\n"
        fname = raw_input(msg).strip()

    quotes = ["'", '"']
    if fname and fname[0] in quotes:
        fname = fname[1:]
    if fname and fname[-1] in quotes:
        fname = fname[:-1]
    if fname:
        try:
            ## root = ivy.tree.read(fname)
            ## IPython.ipapi.get().to_user_ns({treename:root})
            cmd = "%s = ivy.tree.read('%s')" % (treename, fname)
            get_ipython().ex(cmd)
            print "Tree parsed and assigned to variable '%s'" % treename
        except:
            print "Unable to parse tree file '%s'" % fname
    else:
        print "Cancelled"

def __node_completer(self, event):
    symbol = event.symbol
    s = event.line
    if symbol:
        s = s[:-len(symbol)]
    quote = ""
    if s and s[-1] in ["'", '"']:
        quote = s[-1]
        s = s[:-1]
    #base = (re.findall(r'(\w+)\[\Z', s) or [None])[-1]
    base = "".join((re.findall(r'(\w+\.\w*)?(\.)?(\w+)\[\Z', s) or [""])[-1])
    ## print "symbol:", symbol
    ## print "line:", event.line
    ## print "s:", s
    ## print "quote:", quote
    ## print "base:", base
    ## print "obj:", self._ofind(base).get("obj")

    obj = None
    if base:
        obj = self._ofind(base).get("obj")
    ## print '\n'
    ## print 'base', base
    ## print 'obj', obj
    if obj and isinstance(obj, ivy.tree.Node):
        completions = ["'"]
        if quote:
            completions = sorted([ x.label for x in obj.labeled() ])
        return completions

    raise IPython.core.error.TryNext()

try:
    ## import IPython
    IP = get_ipython() #IPython.ipapi.get()
    IP.magic('matplotlib')
    if IP:
        #IP.expose_magic("maketree", __maketree)
        IP.define_magic("maketree", __maketree)
        ## IP.set_hook(
        ##     "complete_command", __node_completer, re_key=r'\[*'
        ##     )
        IP.set_hook(
            "complete_command", __node_completer,
            re_key='.+[[]([\']|["])*\w*$'
            )

except:
    print sys.exc_info()[0]
    sys.stderr.write("Magic commands and completers requires IPython >= 0.11\n")

## if __name__ == "__main__":
##     if len(sys.argv) > 1:
##         for fname in sys.argv[1:]:
##             if os.path.isfile(fname):
##                 execfile(fname)

"""
layout nodes in 2d space

The function of interest is `calc_node_positions` (aka nodepos)
"""
import numpy

class Coordinates:
    """
    Coordinates class for storing xy coordinates
    """
    def __init__(self, x=0, y=0):
        self.x = x
        self.y = y

    def __repr__(self):
        return "Coordinates(%g, %g)" % (self.x, self.y)

    def point(self):
        return (self.x, self.y)

def smooth_xpos(node, n2coords):
    """
    RR: What does smoothing do? -CZ
    """
    if not node.isleaf:
        children = node.children
        for ch in children:
            smooth_xpos(ch, n2coords)

        if node.parent:
            px = n2coords[node.parent].x
            cx = min([ n2coords[ch].x for ch in children ])
            n2coords[node].x = (px + cx)/2.0

def depth_length_preorder_traversal(node, n2coords=None, isroot=False):
    """
    Calculate node depth (root = depth 0) and length to root

    Args:
       node (Node): A node object

    Returns:
        dict: Mapping of nodes to coordinate objects. Coordinate
             objects have attributes "depth" and "length_to_root"
    """
    if n2coords is None:
        n2coords = {}
    coords = n2coords.get(node) or Coordinates()
    coords.node = node
    if (not node.parent) or isroot:
        coords.depth = 0
        coords.length_to_root = 0.0
    else:
        try:
            p = n2coords[node.parent]
            coords.depth = p.depth + 1
            coords.length_to_root = p.length_to_root + (node.length or 0.0)
        except KeyError:
            print node.label, node.parent.label
        except AttributeError:
            coords.depth = 0
            coords.length_to_root = 0
    n2coords[node] = coords

    for ch in node.children:
        depth_length_preorder_traversal(ch, n2coords, False)

    return n2coords

def calc_node_positions(node, width, height,
                        lpad=0, rpad=0, tpad=0, bpad=0,
                        scaled=True, smooth=True, n2coords=None):
    """
    Calculate where nodes should be positioned in 2d space for drawing a tree

    Args:
        node (Node): A (root) node
        width (float): The width of the canvas
        height (float): The height of the canvas
        lpad, rpad, tpad, bpad (float): Padding on the edges of the canvas.
                                Optional, defaults to 0.
        scaled (bool): Whether or not the tree is scaled. Optional, defaults to
                True.
        smooth (bool): Whether or not to smooth the tree. Optional, defaults to
                True.
    Returns:
        dict: Mapping of nodes to Coordinates object
    Notes:
        Origin is at upper left
    """
    width -= (lpad + rpad)
    height -= (tpad + bpad)

    if n2coords is None:
        n2coords = {}
    depth_length_preorder_traversal(node, n2coords=n2coords)
    leaves = node.leaves()
    nleaves = len(leaves)
    maxdepth = max([ n2coords[lf].depth for lf in leaves ])
    unitwidth = width/float(maxdepth)
    unitheight = height/(nleaves-1.0)

    xoff = (unitwidth * 0.5)
    yoff = (unitheight * 0.5)

    if scaled:
        maxlen = max([ n2coords[lf].length_to_root for lf in leaves ])
        scale = width/maxlen

    for i, lf in enumerate(leaves):
        c = n2coords[lf]
        c.y = i * unitheight
        if not scaled:
            c.x = width
        else:
            c.x = c.length_to_root * scale

    for n in node.postiter():
        c = n2coords[n]
        if (not n.isleaf) and n.children:
            children = n.children
            ymax = n2coords[children[0]].y
            ymin = n2coords[children[-1]].y
            c.y = (ymax + ymin)/2.0
            if not scaled:
                c.x = min([ n2coords[ch].x for ch in children ]) - unitwidth
            else:
                c.x = c.length_to_root * scale

    if (not scaled) and smooth:
        for i in range(10):
            smooth_xpos(node, n2coords)

    for coords in n2coords.values():
        coords.x += lpad
        coords.y += tpad

    for n, coords in n2coords.items():
        if n.parent:
            p = n2coords[n.parent]
            coords.px = p.x; coords.py = p.y
        else:
            coords.px = None; coords.py = None

    return n2coords

nodepos = calc_node_positions

def cartesian(node, xscale=1.0, leafspace=None, scaled=True, n2coords=None,
              smooth=0, array=numpy.array, ones=numpy.ones, yunit=None):
    """
    RR: What is the difference between this function and calc_node_positions?
        Is it being used anywhere? -CZ
    """

    if n2coords is None:
        n2coords = {}

    depth_length_preorder_traversal(node, n2coords, True)
    leaves = node.leaves()
    nleaves = len(leaves)

    # leafspace is a vector that should sum to nleaves
    if leafspace is None:
        try: leafspace = [ float(x.leafspace) for x in leaves ]
        except: leafspace = numpy.zeros((nleaves,))
    assert len(leafspace) == nleaves
    #leafspace = array(leafspace)/(sum(leafspace)/float(nleaves))

    maxdepth = max([ n2coords[lf].depth for lf in leaves ])
    depth = maxdepth * xscale
    #if not yunit: yunit = 1.0/nleaves
    yunit = 1.0

    if scaled:
        maxlen = max([ n2coords[lf].length_to_root for lf in leaves ])
        depth = maxlen

    y = 0
    for i, lf in enumerate(leaves):
        c = n2coords[lf]
        yoff = 1 + (leafspace[i] * yunit)
        c.y = y + yoff*0.5
        y += yoff
        if not scaled:
            c.x = depth
        else:
            c.x = c.length_to_root

    for n in node.postiter():
        c = n2coords[n]
        if not n.isleaf:
            children = n.children
            v = [n2coords[children[0]].y, n2coords[children[-1]].y]
            v.sort()
            ymin, ymax = v
            c.y = (ymax + ymin)/2.0
            if not scaled:
                c.x = min([ n2coords[ch].x for ch in children ]) - 1.0
            else:
                c.x = c.length_to_root

    if not scaled:
        for i in range(smooth):
            smooth_xpos(node, n2coords)

    return n2coords

if __name__ == "__main__":
    import tree
    node = tree.read("(a:3,(b:2,(c:4,d:5):1,(e:3,(f:1,g:1):2):2):2);")
    for i, n in enumerate(node.iternodes()):
        if not n.isleaf:
            n.label = "node%s" % i
    node.label = "root"
    n2c = calc_node_positions(node, width=10, height=10, scaled=True)

    from pprint import pprint
    pprint(n2c)

"""
Compute lineages through time
"""
import numpy

# RR: Should results be set to None and then defined in the function to avoid
# problems with mutable defaults in functions? -CZ
def traverse(node, t=0, results=[]):
    """
    Recursively traverse the tree and collect information about when
    nodes split and how many lineages are added by its splitting.
    """
    if node.children:
        ## if not node.label:
        ##     node.label = str(node.id)
        results.append((t, len(node.children)-1))
        for child in node.children:
            traverse(child, t+child.length, results)
    return results

def ltt(node):
    """
    Calculate lineages through time.  The tree is assumed to be an
    ultrametric chronogram (extant leaves, with branch lengths
    proportional to time).

    Args:
        node (Node): A node object. All nodes should have branch lengths.

    Returns:
        tuple: (times, diversity) - 1D-arrays containing the results.
    """
    v = traverse(node) # v is a list of (time, diversity) values
    v.sort()
    # for plotting, it is easiest if x and y values are in separate
    # sequences, so we create a transposed array from v
    times, diversity = numpy.array(v).transpose()
    return times, diversity.cumsum()

def test():
    import newick, ascii
    n = newick.parse("(((a:1,b:2):3,(c:3,d:1):1,(e:0.5,f:3):2.5):1,g:4);")
    v = ltt(n)
    print ascii.render(n, scaled=1)
    for t, n in v:
        print t, n

if __name__ == "__main__":
    test()

"""
Parse newick strings.

The function of interest is `parse`, which returns the root node of
the parsed tree.
"""
import string, sys, re, shlex, types, itertools
import numpy
import nexus
from cStringIO import StringIO
from pprint import pprint

## def read(s):
##     try:
##         s = file(s).read()
##     except:
##         try:
##             s = s.read()
##         except:
##             pass
##     return parse(s)

LABELCHARS = '-.|/?#&'
META = re.compile(r'([^,=\s]+)\s*=\s*(\{[^=}]*\}|"[^"]*"|[^,]+)?')

def add_label_chars(chars):
    global LABELCHARS
    LABELCHARS += chars

class Tokenizer(shlex.shlex):
    """Provides tokens for parsing newick strings."""
    def __init__(self, infile):
        global LABELCHARS
        shlex.shlex.__init__(self, infile, posix=False)
        self.commenters = ''
        self.wordchars = self.wordchars+LABELCHARS
        self.quotes = "'"

    def parse_embedded_comment(self):
        ws = self.whitespace
        self.whitespace = ""
        v = []
        while 1:
            token = self.get_token()
            if token == '':
                sys.stdout.write('EOF encountered mid-comment!\n')
                break
            elif token == ']':
                break
            elif token == '[':
                self.parse_embedded_comment()
            else:
                v.append(token)
        self.whitespace = ws
        return "".join(v)
        ## print "comment:", v

def parse(data, ttable=None, treename=None):
    """
    Parse a newick string.

    Args:
        data: Any file-like object that can be coerced into shlex, or
          a string (converted to StringIO)
        ttable (dict): Mapping of node labels in the newick string
          to other values.

    Returns:
        Node: The root node.
    """
    from tree import Node

    if type(data) in types.StringTypes:
        data = StringIO(data)

    start_pos = data.tell()
    tokens = Tokenizer(data)

    node = None; root = None
    lp=0; rp=0; rooted=1

    previous = None

    ni = 0 # node id counter (preorder) - zero-based indexing
    li = 0 # leaf index counter
    ii = 0 # internal node index counter
    pi = 0 # postorder sequence
    while 1:
        token = tokens.get_token()
        #print token,
        if token == ';' or token == tokens.eof:
            assert lp == rp, \
                   "unbalanced parentheses in tree description: (%s, %s)" \
                   % (lp, rp)
            break

        # internal node
        elif token == '(':
            lp = lp+1
            newnode = Node()
            newnode.ni = ni; ni += 1
            newnode.isleaf = False
            newnode.ii = ii; ii += 1
            newnode.treename = treename
            if node:
                if node.children: newnode.left = node.children[-1].right+1
                else: newnode.left = node.left+1
                node.add_child(newnode)
            else:
                newnode.left = 1; newnode.right = 2
            newnode.right = newnode.left+1
            node = newnode

        elif token == ')':
            rp = rp+1
            node = node.parent
            node.pi = pi; pi += 1
            if node.children:
                node.right = node.children[-1].right + 1

        elif token == ',':
            node = node.parent
            if node.children:
                node.right = node.children[-1].right + 1

        # branch length
        elif token == ':':
            token = tokens.get_token()
            if token == '[':
                node.length_comment = tokens.parse_embedded_comment()
                token = tokens.get_token()

            if not (token == ''):
                try: brlen = float(token)
                except ValueError:
                    raise ValueError, ("invalid literal for branch length, "
                                       "'%s'" % token)
            else:
                raise 'NewickError', \
                      'unexpected end-of-file (expecting branch length)'

            node.length = brlen
        # comment
        elif token == '[':
            node.comment = tokens.parse_embedded_comment()
            if node.comment[0] == '&':
                # metadata
                meta = META.findall(node.comment[1:])
                if meta:
                    for k, v in meta:
                        v = eval(v.replace('{','(').replace('}',')'))
                        node.meta[k] = v

        # leaf node or internal node label
        else:
            if previous != ')': # leaf node
                if ttable:
                    try:
                        ttoken = (ttable.get(int(token)) or
                                  ttable.get(token))
                    except ValueError:
                        ttoken = ttable.get(token)
                    if ttoken:
                        token = ttoken
                newnode = Node()
                newnode.ni = ni; ni += 1
                newnode.pi = pi; pi += 1
                newnode.label = "_".join(token.split()).replace("'", "")
                newnode.isleaf = True
                newnode.li = li; li += 1
                if node.children: newnode.left = node.children[-1].right+1
                else: newnode.left = node.left+1
                newnode.right = newnode.left+1
                newnode.treename = treename
                node.add_child(newnode)
                node = newnode
            else: # label
                if ttable:
                    node.label = ttable.get(token, token)
                else:
                    node.label = token

        previous = token
    node.isroot = True
    return node

## def string(node, length_fmt=":%s", end=True, newline=True):
##     "Recursively create a newick string from node."
##     if not node.isleaf:
##         node_str = "(%s)%s" % \
##                    (",".join([ string(child, length_fmt, False, newline) \
##                                for child in node.children ]),
##                     node.label or ""
##                     )
##     else:
##         node_str = "%s" % node.label

##     if node.length is not None:
##         length_str = length_fmt % node.length
##     else:
##         length_str = ""

##     semicolon = ""
##     if end:
##         if not newline:
##             semicolon = ";"
##         else:
##             semicolon = ";\n"
##     s = "%s%s%s" % (node_str, length_str, semicolon)
##     return s

## def from_nexus(infile, bufsize=None):
##     bufsize = bufsize or 1024*5000
##     TTABLE = re.compile(r'\btranslate\s+([^;]+);', re.I | re.M)
##     TREE = re.compile(r'\btree\s+([_.\w]+)\s*=[^(]+(\([^;]+;)', re.I | re.M)
##     s = infile.read(bufsize)
##     ttable = TTABLE.findall(s) or None
##     if ttable:
##         items = [ shlex.split(line) for line in ttable[0].split(",") ]
##         ttable = dict([ (k, v.replace(" ", "_")) for k, v in items ])
##     trees = TREE.findall(s)
##     ## for i, t in enumerate(trees):
##     ##     t = list(t)
##     ##     if ttable:
##     ##         t[1] = "".join(
##     ##             [ ttable.get(x, "_".join(x.split()).replace("'", ""))
##     ##               for x in shlex.shlex(t[1]) ]
##     ##             )
##     ##     trees[i] = t
##     ## return trees
##     return ttable, trees

def parse_ampersand_comment(s):
    import pyparsing
    pyparsing.ParserElement.enablePackrat()
    from pyparsing import Word, Literal, QuotedString, CaselessKeyword, \
         OneOrMore, Group, Optional, Suppress, Regex, Dict
    word = Word(string.letters+string.digits+"%_")
    key = word.setResultsName("key") + Suppress("=")
    single_value = (Word(string.letters+string.digits+"-.") |
                    QuotedString("'") |
                    QuotedString('"'))
    range_value = Group(Suppress("{") +
                        single_value.setResultsName("min") +
                        Suppress(",") +
                        single_value.setResultsName("max") +
                        Suppress("}"))
    pair = (key + (single_value | range_value).setResultsName("value"))
    g = OneOrMore(pair)
    d = []
    for x in g.searchString(s):
        v = x.value
        if type(v) == str:
            try: v = float(v)
            except ValueError: pass
        else:
            try: v = map(float, v.asList())
            except ValueError: pass
        d.append((x.key, v))
    return d

def nexus_iter(infile):
    import pyparsing
    pyparsing.ParserElement.enablePackrat()
    from pyparsing import Word, Literal, QuotedString, CaselessKeyword, \
         OneOrMore, Group, Optional, Suppress, Regex, Dict
    ## beginblock = Suppress(CaselessKeyword("begin") +
    ##                       CaselessKeyword("trees") + ";")
    ## endblock = Suppress((CaselessKeyword("end") |
    ##                      CaselessKeyword("endblock")) + ";")
    comment = Optional(Suppress("[&") + Regex(r'[^]]+') + Suppress("]"))
    ## translate = CaselessKeyword("translate").suppress()
    name = Word(string.letters+string.digits+"_.") | QuotedString("'")
    ## ttrec = Group(Word(string.digits).setResultsName("number") +
    ##               name.setResultsName("name") +
    ##               Optional(",").suppress())
    ## ttable = Group(translate + OneOrMore(ttrec) + Suppress(";"))
    newick = Regex(r'[^;]+;')
    tree = (CaselessKeyword("tree").suppress() +
            Optional("*").suppress() +
            name.setResultsName("tree_name") +
            comment.setResultsName("tree_comment") +
            Suppress("=") +
            comment.setResultsName("root_comment") +
            newick.setResultsName("newick"))
    ## treesblock = Group(beginblock +
    ##                    Optional(ttable.setResultsName("ttable")) +
    ##                    Group(OneOrMore(tree)) +
    ##                    endblock)

    def not_begin(s): return s.strip().lower() != "begin trees;"
    def not_end(s): return s.strip().lower() not in ("end;", "endblock;")
    def parse_ttable(f):
        ttable = {}
        while True:
            s = f.next().strip()
            if not s: continue
            if s.lower() == ";": break
            if s[-1] == ",": s = s[:-1]
            k, v = s.split()
            ttable[k] = v
            if s[-1] == ";": break
        return ttable

    # read lines between "begin trees;" and "end;"
    f = itertools.takewhile(not_end, itertools.dropwhile(not_begin, infile))
    s = f.next().strip().lower()
    if s != "begin trees;":
        print sys.stderr, "Expecting 'begin trees;', got %s" % s
        raise StopIteration
    ttable = {}
    while True:
        try: s = f.next().strip()
        except StopIteration: break
        if not s: continue
        if s.lower() == "translate":
            ttable = parse_ttable(f)
            print "ttable: %s" % len(ttable)
        elif s.split()[0].lower()=='tree':
            match = tree.parseString(s)
            yield nexus.Newick(match, ttable)

## def test():
##     with open("/home/rree/Dropbox/pedic-comm-amnat/phylo/beast-results/"
##               "simple_stigma.trees.log") as f:
##         for rec in nexus_iter(f):
##             r = parse(rec.newick, ttable=rec.ttable)
##             for x in r: print x, x.comments

def test_parse_comment():
    v = (("height_median=1.1368683772161603E-13,height=9.188229043880098E-14,"
          "height_95%_HPD={5.6843418860808015E-14,1.7053025658242404E-13},"
          "height_range={0.0,2.8421709430404007E-13}"),
         "R", "lnP=-154.27154502342688,lnP=-24657.14341301901",
         'states="T-lateral"')
    for s in v:
        print "input:", s
        print dict(parse_ampersand_comment(s))

import itertools
from collections import defaultdict
import newick

class Newick(object):
    """
    convenience class for storing the results of a newick tree
    record from a nexus file, as parsed by newick.nexus_iter
    """
    def __init__(self, parse_results=None, ttable={}):
        self.name = ""
        self.comment = ""
        self.root_comment = ""
        self.newick = ""
        self.ttable = ttable
        if parse_results: self.populate(parse_results)

    def populate(self, parse_results, ttable={}):
        self.name = parse_results.tree_name
        self.comment = parse_results.tree_comment
        self.root_comment = parse_results.root_comment
        self.newick = parse_results.newick
        if ttable: self.ttable = ttable

    def parse(self, newick=newick):
        assert self.newick
        self.root = newick.parse(
            self.newick, ttable=self.ttable, treename=self.name
            )
        return self.root

def fetchaln(fname):
    """Fetch alignment"""
    from Bio.Nexus import Nexus
    n = Nexus.Nexus(fname)
    return n

def split_blocks(infile):
    from cStringIO import StringIO
    dropwhile = itertools.dropwhile; takewhile = itertools.takewhile
    blocks = []
    not_begin = lambda s: not s.lower().startswith("begin")
    not_end = lambda s: not s.strip().lower() in ("end;", "endblock;")
    while 1:
        f = takewhile(not_end, dropwhile(not_begin, infile))
        try:
            b = f.next().split()[-1][:-1]
            blocks.append((b, StringIO("".join(list(f)))))
        except StopIteration:
            break
    return blocks

def parse_treesblock(infile):
    import string
    from pyparsing import Optional, Word, Regex, CaselessKeyword, Suppress
    from pyparsing import QuotedString
    comment = Optional(Suppress("[&") + Regex(r'[^]]+') + Suppress("]"))
    name = Word(string.letters+string.digits+"_") | QuotedString("'")
    newick = Regex(r'[^;]+;')
    tree = (CaselessKeyword("tree").suppress() +
            Optional("*").suppress() +
            name.setResultsName("tree_name") +
            comment.setResultsName("tree_comment") +
            Suppress("=") +
            comment.setResultsName("root_comment") +
            newick.setResultsName("newick"))
    ## treesblock = Group(beginblock +
    ##                    Optional(ttable.setResultsName("ttable")) +
    ##                    Group(OneOrMore(tree)) +
    ##                    endblock)

    def parse_ttable(f):
        ttable = {}
        while True:
            s = f.next().strip()
            if s.lower() == ";": break
            if s[-1] in ",;": s = s[:-1]
            k, v = s.split()
            ttable[k] = v
            if s[-1] == ";": break
        return ttable

    ttable = {}
    while True:
        try: s = infile.next().strip()
        except StopIteration: break
        if s.lower() == "translate":
            ttable = parse_ttable(infile)
            print "ttable: %s" % len(ttable)
        else:
            match = tree.parseString(s)
            yield Newick(match, ttable)

from itertools import izip, imap
import numpy

def finditer(seq, substr, start=0):
    """
    Find substrings within a sequence

    Args:
        seq (str): A sequence.
        substr (str): A subsequence to search for
        start (int): Starting index. Defaults to 0
    Yields:
        int: Starting indicies of where the substr was found in seq
    """
    N = len(substr)
    i = seq.find(substr, start)
    while i >= 0:
        yield i
        i = seq.find(substr, i+N)

def gapidx(seq, gapchar='-'):
    """
    For a sequence with gaps, calculate site positions without gaps

    Args:
        seq (list): Each element of the list is one character in a sequence.
        gapchar (str): The character gaps are coded as. Defaults to '-'
    Returns:
        array: An array where the first element corresponds to range(number of
        characters that are not gaps) and the second element is the indicies
        of all characters that are not gaps.
    """
    a = numpy.array(seq)
    idx = numpy.arange(len(a))
    nongap = idx[a != gapchar]
    return numpy.array((numpy.arange(len(nongap)), nongap))

def find_stop_codons(seq, pos=0):
    """
    Find stop codons within sequence (in reading frame)

    Args:
        seq (str): A sequence
        pos (int): Starting position. Defaults to 0.
    Yields:
        tuple: The index where the stop codon starts
        and which stop codon was found.
    """
    s = seq[pos:]
    it = iter(s)
    g = imap(lambda x:"".join(x), izip(it, it, it))
    for i, x in enumerate(g):
        if x in ("TAG", "TAA", "TGA"):
            yield pos+(i*3), x

from operator import itemgetter
from heapq import nlargest
from itertools import repeat, ifilter

class Storage(dict):
    """
    A Storage object is like a dictionary except `obj.foo` can be used
    in addition to `obj['foo']`.

    From web2py/gluon/storage.py by Massimo Di Pierro (www.web2py.com)
    """
    def __getattr__(self, key):
        try:
            return self[key]
        except KeyError:
            return None

    def __setattr__(self, key, value):
        self[key] = value

    def __delattr__(self, key):
        try:
            del self[key]
        except KeyError, k:
            raise AttributeError, k

    def __repr__(self):
        return '<Storage ' + dict.__repr__(self) + '>'

    def __getstate__(self):
        return dict(self)

    def __setstate__(self, value):
        for (k, v) in value.items():
            self[k] = v

class MaxDict(dict):
    def __setitem__(self, key, value):
        v = self.get(key)
        if value > v:
            dict.__setitem__(self, key, value)

#from http://code.activestate.com/recipes/576611/
class Counter(dict):
    """Dict subclass for counting hashable objects.  Sometimes called a bag
    or multiset.  Elements are stored as dictionary keys and their counts
    are stored as dictionary values.

    >>> Counter('zyzygy')
    Counter({'y': 3, 'z': 2, 'g': 1})

    """

    def __init__(self, iterable=None, **kwds):
        """Create a new, empty Counter object.  And if given, count elements
        from an input iterable.  Or, initialize the count from another mapping
        of elements to their counts.

        >>> c = Counter()                           # a new, empty counter
        >>> c = Counter('gallahad')                 # a new counter from an iterable
        >>> c = Counter({'a': 4, 'b': 2})           # a new counter from a mapping
        >>> c = Counter(a=4, b=2)                   # a new counter from keyword args

        """
        self.update(iterable, **kwds)

    def __missing__(self, key):
        return 0

    def most_common(self, n=None):
        """List the n most common elements and their counts from the most
        common to the least.  If n is None, then list all element counts.

        >>> Counter('abracadabra').most_common(3)
        [('a', 5), ('r', 2), ('b', 2)]

        """
        if n is None:
            return sorted(self.iteritems(), key=itemgetter(1), reverse=True)
        return nlargest(n, self.iteritems(), key=itemgetter(1))

    def elements(self):
        """Iterator over elements repeating each as many times as its count.

        >>> c = Counter('ABCABC')
        >>> sorted(c.elements())
        ['A', 'A', 'B', 'B', 'C', 'C']

        If an element's count has been set to zero or is a negative number,
        elements() will ignore it.

        """
        for elem, count in self.iteritems():
            for _ in repeat(None, count):
                yield elem

    # Override dict methods where the meaning changes for Counter objects.

    @classmethod
    def fromkeys(cls, iterable, v=None):
        raise NotImplementedError(
            'Counter.fromkeys() is undefined.  Use Counter(iterable) instead.')

    def update(self, iterable=None, **kwds):
        """Like dict.update() but add counts instead of replacing them.

        Source can be an iterable, a dictionary, or another Counter instance.

        >>> c = Counter('which')
        >>> c.update('witch')           # add elements from another iterable
        >>> d = Counter('watch')
        >>> c.update(d)                 # add elements from another counter
        >>> c['h']                      # four 'h' in which, witch, and watch
        4

        """
        if iterable is not None:
            if hasattr(iterable, 'iteritems'):
                if self:
                    self_get = self.get
                    for elem, count in iterable.iteritems():
                        self[elem] = self_get(elem, 0) + count
                else:
                    dict.update(self, iterable) # fast path when counter is empty
            else:
                self_get = self.get
                for elem in iterable:
                    self[elem] = self_get(elem, 0) + 1
        if kwds:
            self.update(kwds)

    def copy(self):
        'Like dict.copy() but returns a Counter instance instead of a dict.'
        return Counter(self)

    def __delitem__(self, elem):
        'Like dict.__delitem__() but does not raise KeyError for missing values.'
        if elem in self:
            dict.__delitem__(self, elem)

    def __repr__(self):
        if not self:
            return '%s()' % self.__class__.__name__
        items = ', '.join(map('%r: %r'.__mod__, self.most_common()))
        return '%s({%s})' % (self.__class__.__name__, items)

    # Multiset-style mathematical operations discussed in:
    #       Knuth TAOCP Volume II section 4.6.3 exercise 19
    #       and at http://en.wikipedia.org/wiki/Multiset
    #
    # Outputs guaranteed to only include positive counts.
    #
    # To strip negative and zero counts, add-in an empty counter:
    #       c += Counter()

    def __add__(self, other):
        """Add counts from two counters.

        >>> Counter('abbb') + Counter('bcc')
        Counter({'b': 4, 'c': 2, 'a': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        result = Counter()
        for elem in set(self) | set(other):
            newcount = self[elem] + other[elem]
            if newcount > 0:
                result[elem] = newcount
        return result

    def __sub__(self, other):
        """ Subtract count, but keep only results with positive counts.

        >>> Counter('abbbc') - Counter('bccd')
        Counter({'b': 2, 'a': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        result = Counter()
        for elem in set(self) | set(other):
            newcount = self[elem] - other[elem]
            if newcount > 0:
                result[elem] = newcount
        return result

    def __or__(self, other):
        """Union is the maximum of value in either of the input counters.

        >>> Counter('abbb') | Counter('bcc')
        Counter({'b': 3, 'c': 2, 'a': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        _max = max
        result = Counter()
        for elem in set(self) | set(other):
            newcount = _max(self[elem], other[elem])
            if newcount > 0:
                result[elem] = newcount
        return result

    def __and__(self, other):
        """ Intersection is the minimum of corresponding counts.

        >>> Counter('abbb') & Counter('bcc')
        Counter({'b': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        _min = min
        result = Counter()
        if len(self) < len(other):
            self, other = other, self
        for elem in ifilter(self.__contains__, other):
            newcount = _min(self[elem], other[elem])
            if newcount > 0:
                result[elem] = newcount
        return result

def convert(d):
    "convert a (potentially nested) dict to Storage"
    from types import DictType
    t = type(d)
    if t == DictType:
        for k,v in d.items():
            d[k] = convert(v)
        return Storage(d)
    return d

if __name__ == '__main__':
    import doctest
    print doctest.testmod()

"""
The Node class and functions for creating trees from Newick strings,
etc.

ivy does not have a Tree class per se, as most functions operate
directly on Node objects.
"""
import os, types
from storage import Storage
from copy import copy as _copy
from matrix import vcv
import newick
from itertools import izip_longest

## class Tree(object):
##     """
##     A simple Tree class.
##     """
##     def __init__(self, data=None, format="newick", name=None, ttable=None):
##         self.root = None
##         if data:
##             self.root = read(data, format, name, ttable)
##         self.name = name
##         self.ttable = ttable

##     def __getattribute__(self, a):
##         r = object.__getattribute__(self, 'root')
##         try:
##             return object.__getattribute__(r, a)
##         except AttributeError:
##             return object.__getattribute__(self, a)

def traverse(node):
    "recursive preorder iterator based solely on .children attribute"
    yield node
    for child in node.children:
        for descendant in traverse(child):
            yield descendant

class Node(object):
    """
    A basic Node class with attributes and references to child nodes
    ('children', a list) and 'parent'.

    Keyword Args:
        id: ID of the node. If not provided, is set using
          builtin id function
        ni (int): Node index.
        li (int): Leaf index.
        isroot (bool): Is the node a root.
        isleaf (bool): Is the node a leaf.
        label (str): Node label.
        length (float): Branch length from node to parent
        support: RR: Are these bootstrap support values? -CZ
        age (float): Age of the node in time units.
        parent (Node): Parent of the ndoe.
        children (list): List of node objects. Children of node
        nchildren (int): No. of children
        left: RR: Unsure what left and right mean -CZ
        treename: Name of tree
        comment: Comments for tree

    """
    def __init__(self, **kwargs):
        self.id = None
        self.ni = None # node index
        self.li = None # leaf index
        self.isroot = False
        self.isleaf = False
        self.label = None
        self.length = None
        self.support = None
        self.age = None
        self.parent = None
        self.children = []
        self.nchildren = 0
        self.left = None
        self.right = None
        self.treename = ""
        self.comment = ""
        self.meta = {}
        ## self.length_comment = ""
        ## self.label_comment = ""
        if kwargs:
            for k, v in kwargs.items():
                setattr(self, k, v)
        if self.id is None: self.id = id(self)

    def __copy__(self):
        return self.copy()

    def __repr__(self):
        v = []
        if self.isroot:
            v.append("root")
        elif self.isleaf:
            v.append("leaf")

        if self.label:
            v.append("'%s'" % self.label)

        s = ", ".join(v)

        nid = ((self.id if (self.id is not None) else self.ni) or
               '<%s>' % id(self))
        if s:
            s = "Node(%s, %s)" % (nid, s)
        else:
            s = "Node(%s)" % nid
        return s

    def __contains__(self, other):
        """
        For use with `in` keyword

        Args:
            other: Another node or node label.
        Returns:
            bool: Whether or not the other node is a descendant of self
        """
        otype = type(other)
        if other and otype in types.StringTypes:
            for x in self:
                if other == x.label:
                    return True
            return False
        else:
            assert otype == type(self)
            for x in self.iternodes():
                if other == x:
                    return True
        return False

    def __iter__(self):
        for node in self.iternodes():
            yield node

    def __len__(self):
        """
        Number of nodes descended from self

        Returns:
            int: Number of nodes descended from self (including self)
        """
        i = 0
        for n in self:
            i += 1
        return i

    def __nonzero__(self):
        return True

    def __getitem__(self, x):
        """
        Args:
            x: A Node, Node.id (int) or a Node.label (string)

        Returns:
            Node: Found node(s)

        """
        for n in self:
            if n==x or n.id==x or n.ni == x or (n.label and n.label==x):
                return n
        raise IndexError(str(x))

    def ascii(self, *args, **kwargs):
        """
        Create ascii tree.

        Keyword Args:
            unitlen (float): How long each unit should be rendered as.
              Defaults to 3.
            minwidth (float): Minimum width of the plot. Defaults to 50
            maxwidth (float): Maximum width of the plot. Defaults to None
            scaled (bool): Whether or not the tree is scaled. Defaults to False
            show_internal_labels (bool): Whether or not to show labels
              on internal nodes. Defaults to True.
        Returns:
            str: Ascii tree to be shown with print().
        """
        from ascii import render
        return render(self, *args, **kwargs)

    def collapse(self, add=False):
        """
        Remove self and collapse children to polytomy

        Args:
            add (bool): Whether or not to add self's length to children's
              length.

        Returns:
            Node: Parent of self

        """
        assert self.parent
        p = self.prune()
        for c in self.children:
            p.add_child(c)
            if add and (c.length is not None):
                c.length += self.length
        self.children = []
        return p

    def copy(self, recurse=False):
        """
        Return a copy of the node, but not copies of children, parent,
        or any attribute that is a Node.

        If `recurse` is True, recursively copy child nodes.

        Args:
            recurse (bool): Whether or not to copy children as well as self.

        Returns:
            Node: A copy of self.

        TODO: test this function.

        RR: This function runs rather slowly -CZ
        """
        newnode = Node()
        for attr, value in self.__dict__.items():
            if (attr not in ("children", "parent") and
                not isinstance(value, Node)):
                setattr(newnode, attr, _copy(value))
            if recurse:
                newnode.children = [
                    child.copy(True) for child in self.children
                    ]
        return newnode

    def leafsets(self, d=None, labels=False):
        """return a mapping of nodes to leaf sets (nodes or labels)"""
        d = d or {}
        if not self.isleaf:
            s = set()
            for child in self.children:
                if child.isleaf:
                    if labels:
                        s.add(child.label)
                    else:
                        s.add(child)
                else:
                    d = child.leafsets(d, labels)
                    s = s | d[child]
            d[self] = frozenset(s)
        return d

    def mrca(self, *nodes):
        """
        Find most recent common ancestor of *nodes*

        Args:
            *nodes (Node): Node objects
        Returns:
            Node: The MRCA of *nodes*
        """
        if len(nodes) == 1:
            nodes = tuple(nodes[0])
        if len(nodes) == 1:
            return nodes[0]
        nodes = set([ self[n] for n in nodes ])
        anc = []
        def f(n):
            seen = set()
            for c in n.children: seen.update(f(c))
            if n in nodes: seen.add(n)
            if seen == nodes and (not anc): anc.append(n)
            return seen
        f(self)
        return anc[0]

    ## def mrca(self, *nodes):
    ##     """
    ##     Find most recent common ancestor of *nodes*
    ##     """
    ##     if len(nodes) == 1:
    ##         nodes = tuple(nodes[0])
    ##     if len(nodes) == 1:
    ##         return nodes[0]
    ##     ## assert len(nodes) > 1, (
    ##     ##     "Need more than one node for mrca(), got %s" % nodes
    ##     ##     )
    ##     def f(x):
    ##         if isinstance(x, Node):
    ##             return x
    ##         elif type(x) in types.StringTypes:
    ##             return self.find(x)
    ##     nodes = map(f, nodes)
    ##     assert all(filter(lambda x: isinstance(x, Node), nodes))

    ##     #v = [ list(n.rootpath()) for n in nodes if n in self ]
    ##     v = [ list(x) for x in izip_longest(*[ reversed(list(n.rootpath()))
    ##                                            for n in nodes if n in self ]) ]
    ##     if len(v) == 1:
    ##         return v[0][0]
    ##     anc = None
    ##     for x in v:
    ##         s = set(x)
    ##         if len(s) == 1: anc = list(s)[0]
    ##         else: break
    ##     return anc

    def ismono(self, *leaves):
        """
        Test if leaf descendants are monophyletic

        Args:
            *leaves (Node): At least two leaf Node objects

        Returns:
            bool: Are the leaf descendants monophyletic?

        RR: Should this function have a check to make sure the input nodes are
        leaves? There is some strange behavior if you input internal nodes -CZ
        """
        if len(leaves) == 1:
            leaves = list(leaves)[0]
        assert len(leaves) > 1, (
            "Need more than one leaf for ismono(), got %s" % leaves
            )
        anc = self.mrca(leaves)
        if anc:
            return bool(len(anc.leaves())==len(leaves))

    def order_subtrees_by_size(self, n2s=None, recurse=False, reverse=False):
        """
        Order interal clades by size

        """
        if n2s is None:
            n2s = clade_sizes(self)
        if not self.isleaf:
            v = [ (n2s[c], c.label, c) for c in self.children ]
            v.sort()
            if reverse:
                v.reverse()
            self.children = [ x[-1] for x in v ]
            if recurse:
                for c in self.children:
                    c.order_subtrees_by_size(n2s, recurse=True, reverse=reverse)

    def ladderize(self, reverse=False):
        self.order_subtrees_by_size(recurse=True, reverse=reverse)
        return self

    def add_child(self, child):
        """
        Add child as child of self

        Args:
            child (Node): A node object

        """
        assert child not in self.children
        self.children.append(child)
        child.parent = self
        child.isroot = False
        self.nchildren += 1

    def bisect_branch(self):
        """
        Add new node as parent to self in the middle of branch to parent.

        Returns:
            Node: A new node.

        """
        assert self.parent
        parent = self.prune()
        n = Node()
        if self.length:
            n.length = self.length/2.0
            self.length /= 2.0
        parent.add_child(n)
        n.add_child(self)
        return n

    def remove_child(self, child):
        """
        Remove child.

        Args:
            child (Node): A node object that is a child of self

        """
        assert child in self.children
        self.children.remove(child)
        child.parent = None
        self.nchildren -= 1
        if not self.children:
            self.isleaf = True

    def labeled(self):
        """
        Return a list of all descendant nodes that are labeled

        Returns:
            list: All descendants of self that are labeled (including self)
        """
        return [ n for n in self if n.label ]

    def leaves(self, f=None):
        """
        Return a list of leaves. Can be filtered with f.

        Args:
            f (function): A function that evaluates to True if called with desired
              node as the first input

        Returns:
            list: A list of leaves that are true for f (if f is given)

        """
        if f: return [ n for n in self if (n.isleaf and f(n)) ]
        return [ n for n in self if n.isleaf ]

    def internals(self, f=None):
        """
        Return a list nodes that have children (internal nodes)

        Args:
            f (function): A function that evaluates to true if called with desired
              node as the first input

        Returns:
            list: A list of internal nodes that are true for f (if f is given)

        """
        if f: return [ n for n in self if (n.children and f(n)) ]
        return [ n for n in self if n.children ]

    def clades(self):
        """
        Get internal nodes descended from self

        Returns:
            list: A list of internal nodes descended from (and not including) self.

        """
        return [ n for n in self if (n is not self) and not n.isleaf ]

    def iternodes(self, f=None):
        """
        Return a generator of nodes descendant from self - including self

        Args:
            f (function): A function that evaluates to true if called with
            desired node as the first input

        Yields:
            Node: Nodes descended from self (including self) in
              preorder sequence

        """
        if (f and f(self)) or (not f):
            yield self
        for child in self.children:
            for n in child.iternodes(f):
                yield n

    def iterleaves(self):
        """
        Yield leaves descendant from self
        """
        return self.iternodes(lambda x:x.isleaf)

    def preiter(self, f=None):
        """
        Yield nodes in preorder sequence
        """
        for n in self.iternodes(f=f):
            yield n

    def postiter(self, f=None):
        """
        Yield nodes in postorder sequence
        """
        if not self.isleaf:
            for child in self.children:
                for n in child.postiter():
                    if (f and f(n)) or (not f):
                        yield n
        if (f and f(self)) or (not f):
            yield self

    def descendants(self, order="pre", v=None, f=None):
        """
        Return a list of nodes descendant from self - but _not_
        including self!

        Args:
            order (str): Indicates wether to return nodes in preorder or
              postorder sequence. Optional, defaults to "pre"
            f (function): filtering function that evaluates to True if desired
              node is called as the first parameter.

        Returns:
            list: A list of nodes descended from self not including self.

        """
        v = v or []
        for child in self.children:
            if (f and f(child)) or (not f):
                if order == "pre":
                    v.append(child)
                else:
                    v.insert(0, child)
            if child.children:
                child.descendants(order, v, f)
        return v

    def get(self, f, *args, **kwargs):
        """
        Return the first node found by node.find()

        Args:
            f (function): A function that evaluates to True if desired
              node is called as the first parameter.
        Returns:
            Node: The first node found by node.find()

        """
        v = self.find(f, *args, **kwargs)
        try:
            return v.next()
        except StopIteration:
            return None

    def grep(self, s, ignorecase=True):
        """
        Find nodes by regular-expression search of labels

        Args:
            s (str): String to search.
            ignorecase (bool): Indicates to ignore case. Defaults to true.

        Returns:
            lsit: A list of node objects whose labels were matched by s.

        """
        import re
        if ignorecase:
            pattern = re.compile(s, re.IGNORECASE)
        else:
            pattern = re.compile(s)

        search = pattern.search
        return [ x for x in self if x.label and search(x.label) ]

    def lgrep(self, s, ignorecase=True):
        """
        Find leaves by regular-expression search of labels

        Args:
            s (str): String to search.
            ignorecase (bool): Indicates to ignore case. Defaults to true.

        Returns:
            lsit: A list of node objects whose labels were matched by s.

        """
        return [ x for x in self.grep(s, ignorecase=ignorecase) if x.isleaf ]

    def bgrep(self, s, ignorecase=True):
        """
        Find branches (internal nodes) by regular-expression search of
        labels

        Args:
            s (str): String to search.
            ignorecase (bool): Indicates to ignore case. Defaults to true.

        Returns:
            lsit: A list of node objects whose labels were matched by s.

        """
        return [ x for x in self.grep(s, ignorecase=ignorecase) if
               (not x.isleaf) ]

    def find(self, f, *args, **kwargs):
        """
        Find descendant nodes.

        Args:
            f: Function or a string.  If a string, it is converted to a
              function for finding *f* as a substring in node labels.
              Otherwise, *f* should evaluate to True if called with a desired
              node as the first parameter.

        Yields:
            Node: Found nodes in preorder sequence.

        """
        if not f: return
        if type(f) in types.StringTypes:
            func = lambda x: (f or None) in (x.label or "")
        else:
            func = f
        for n in self.iternodes():
            if func(n, *args, **kwargs):
                yield n

    def findall(self, f, *args, **kwargs):
        """Return a list of found nodes."""
        return list(self.find(f, *args, **kwargs))

    def prune(self):
        """
        Remove self if self is not root.

        Returns:
            Node: Parent of self. If parent had only two children,
              parent is now a 'knee' and can be removed with excise.

        """
        p = self.parent
        if p:
            p.remove_child(self)
        return p

    def excise(self):
        """
        For 'knees': remove self from between parent and single child
        """
        assert self.parent
        assert len(self.children)==1
        p = self.parent
        c = self.children[0]
        if c.length is not None and self.length is not None:
            c.length += self.length
        c.prune()
        self.prune()
        p.add_child(c)
        return p

    def graft(self, node):
        """
        Add node as sister to self.
        """
        parent = self.parent
        parent.remove_child(self)
        n = Node()
        n.add_child(self)
        n.add_child(node)
        parent.add_child(n)

    ## def leaf_distances(self, store=None, measure="length"):
    ##     """
    ##     for each internal node, calculate the distance to each leaf,
    ##     measured in branch length or internodes
    ##     """
    ##     if store is None:
    ##         store = {}
    ##     leaf2len = {}
    ##     if self.children:
    ##         for child in self.children:
    ##             if measure == "length":
    ##                 dist = child.length
    ##             elif measure == "nodes":
    ##                 dist = 1
    ##             child.leaf_distances(store, measure)
    ##             if child.isleaf:
    ##                 leaf2len[child] = dist
    ##             else:
    ##                 for k, v in store[child].items():
    ##                     leaf2len[k] = v + dist
    ##     else:
    ##         leaf2len[self] = {self: 0}
    ##     store[self] = leaf2len
    ##     return store

    def leaf_distances(self, measure="length"):
        """
        RR: I don't quite understand the structure of the output. Also,
            I can't figure out what "measure" does.-CZ
        """
        from collections import defaultdict
        store = defaultdict(lambda:defaultdict(lambda:0))
        nodes = [ x for x in self if x.children ]
        for lf in self.leaves():
            x = lf.length
            for n in lf.rootpath(self):
                store[n][lf] = x
                x += (n.length or 0)
        return store

    def rootpath(self, end=None, stop=None):
        """
        Iterate over parent nodes toward the root, or node *end* if
        encountered.

        Args:
            end (Node): A Node object to iterate to (instead of iterating
              towards root). Optional, defaults to None
            stop (function): A function that returns True if desired node is called
              as the first parameter. Optional, defaults to None

        Yields:
            Node: Nodes in path to root (or end).

        """
        n = self.parent
        while 1:
            if n is None: raise StopIteration
            yield n
            if n.isroot or (end and n == end) or (stop and stop(n)):
                raise StopIteration
            n = n.parent

    def rootpath_length(self, end=None):
        """
        Get length from self to root(if end is None) or length
        from self to an ancestor node (if end is an ancestor to self)

        Args:
            end (Node): A node object

        Returns:
            float: The length from self to root/end

        """
        n = self
        x = 0.0
        while n.parent:
            x += n.length
            if n.parent == end:
                break
            n = n.parent
        return x
        ## f = lambda x:x.parent==end
        ## v = [self.length]+[ x.length for x in self.rootpath(stop=f)
        ##                     if x.parent ]
        ## assert None not in v
        ## return sum(v)

    def max_tippath(self, first=True):
        """
        Get the maximum length from self to a leaf node
        """
        v = 0
        if self.children:
            v = max([ c.max_tippath(False) for c in self.children ])
        if not first:
            if self.length is None: v += 1
            else: v += self.length
        return v

    def subtree_mapping(self, labels, clean=False):
        """
        Find the set of nodes in 'labels', and create a new tree
        representing the subtree connecting them.  Nodes are assumed
        to be non-nested.

        Returns:
            dict: a mapping of old nodes to new nodes and vice versa.

        TODO: test this, high bug probability
        """
        d = {}
        oldtips = [ x for x in self.leaves() if x.label in labels ]
        for tip in oldtips:
            path = list(tip.rootpath())
            for node in path:
                if node not in d:
                    newnode = Node()
                    newnode.isleaf = node.isleaf
                    newnode.length = node.length
                    newnode.label = node.label
                    d[node] = newnode
                    d[newnode] = node
                else:
                    newnode = d[node]

                for child in node.children:
                    if child in d:
                        newchild = d[child]
                        if newchild not in newnode.children:
                            newnode.add_child(newchild)
        d["oldroot"] = self
        d["newroot"] = d[self]
        if clean:
            n = d["newroot"]
            while 1:
                if n.nchildren == 1:
                    oldnode = d[n]
                    del d[oldnode]; del d[n]
                    child = n.children[0]
                    child.parent = None
                    child.isroot = True
                    d["newroot"] = child
                    d["oldroot"] = d[child]
                    n = child
                else:
                    break

            for tip in oldtips:
                newnode = d[tip]
                while 1:
                    newnode = newnode.parent
                    oldnode = d[newnode]
                    if newnode.nchildren == 1:
                        child = newnode.children[0]
                        if newnode.length:
                            child.length += newnode.length
                        newnode.remove_child(child)
                        if newnode.parent:
                            parent = newnode.parent
                            parent.remove_child(newnode)
                            parent.add_child(child)
                        del d[oldnode]; del d[newnode]
                    if not newnode.parent:
                        break

        return d

    def reroot_orig(self, newroot):
        assert newroot in self
        self.isroot = False
        newroot.isroot = True
        v = []
        n = newroot
        while 1:
            v.append(n)
            if not n.parent: break
            n = n.parent
        v.reverse()
        for i, cp in enumerate(v[:-1]):
            node = v[i+1]
            # node is current node; cp is current parent
            cp.remove_child(node)
            node.add_child(cp)
            cp.length = node.length
        return newroot

    def reroot(self, newroot):
        """
        RR: I can't get this to work properly -CZ
        """
        newroot = self[newroot]
        assert newroot in self
        self.isroot = False
        n = newroot
        v = list(n.rootpath())
        v.reverse()
        for node in (v+[n])[1:]:
            # node is current node; cp is current parent
            cp = node.parent
            cp.remove_child(node)
            node.add_child(cp)
            cp.length = node.length
            cp.label = node.label
        newroot.isroot = True
        return newroot

    def makeroot(self, shift_labels=False):
        """
        shift_labels: flag to shift internal parent-child node labels
        when internode polarity changes; suitable e.g. if internal node
        labels indicate unrooted bipartition support
        """
        v = list(self.rootpath())
        v[-1].isroot = False
        v.reverse()
        for node in v[1:] + [self]:
            # node is current node; cp is current parent
            cp = node.parent
            cp.remove_child(node)
            node.add_child(cp)
            cp.length = node.length
            if shift_labels:
                cp.label = node.label
        self.isroot = True
        return self

    def write(self, outfile=None, format="newick", length_fmt=":%g", end=True,
              clobber=False):
        if format=="newick":
            s = write_newick(self, outfile, length_fmt, True, clobber)
            if not outfile:
                return s

reroot = Node.reroot

def index(node, n=0, d=0):
    """
    recursively attach 'next', 'back', (and 'left', 'right'), 'ni',
    'ii', 'pi', and 'node_depth' attributes to nodes
    """
    node.next = node.left = n
    if not node.parent:
        node.node_depth = d
    else:
        node.node_depth = node.parent.node_depth + 1
    n += 1
    for i, c in enumerate(node.children):
        if i > 0:
            n = node.children[i-1].back + 1
        index(c, n)

    if node.children:
        node.back = node.right = node.children[-1].back + 1
    else:
        node.back = node.right = n
    return node.back

def remove_singletons(root, add=True):
    "Remove descendant nodes that are the sole child of their parent"
    for leaf in root.leaves():
        for n in leaf.rootpath():
            if n.parent and len(n.parent.children)==1:
                n.collapse(add)

def cls(root):
    """
    Get clade sizes of whole tree
    Args:
        * root: A root node

    Returns:
        * A dict mapping nodes to clade sizes

    """
    results = {}
    for node in root.postiter():
        if node.isleaf:
            results[node] = 1
        else:
            results[node] = sum(results[child] for child in node.children)
    return results

def clade_sizes(node, results={}):
    """Map node and descendants to number of descendant tips"""
    size = int(node.isleaf)
    if not node.isleaf:
        for child in node.children:
            clade_sizes(child, results)
            size += results[child]
    results[node] = size
    return results

def write(node, outfile=None, format="newick", length_fmt=":%g",
          clobber=False):
    if format=="newick" or ((type(outfile) in types.StringTypes) and
                            (outfile.endswith(".newick") or
                             outfile.endswith(".new"))):
        s = write_newick(node, outfile, length_fmt, True, clobber)
        if not outfile:
            return s

def write_newick(node, outfile=None, length_fmt=":%g", end=False,
                 clobber=False):
    if not node.isleaf:
        node_str = "(%s)%s" % \
                   (",".join([ write_newick(child, outfile, length_fmt,
                                            False, clobber)
                               for child in node.children ]),
                    (node.label or "")
                    )
    else:
        node_str = "%s" % node.label

    if node.length is not None:
        length_str = length_fmt % node.length
    else:
        length_str = ""

    semicolon = ""
    if end:
        semicolon = ";"
    s = "%s%s%s" % (node_str, length_str, semicolon)
    if end and outfile:
        flag = False
        if type(outfile) in types.StringTypes:
            if not clobber:
                assert not os.path.isfile(outfile), "File '%s' exists! (Set clobber=True to overwrite)" % outfile
            flag = True
            outfile = file(outfile, "w")
        outfile.write(s)
        if flag:
            outfile.close()
    return s

def read(data, format=None, treename=None, ttable=None):
    """
    Read a single tree from *data*, which can be a Newick string, a
    file name, or a file-like object with `tell` and 'read`
    methods. *treename* is an optional string that will be attached to
    all created nodes.

    Args:
        data: A file or file-like object or newick string

    Returns:
        Node: The root node.
    """
    import newick
    StringTypes = types.StringTypes

    def strip(s):
        fname = os.path.split(s)[-1]
        head, tail = os.path.splitext(fname)
        tail = tail.lower()
        if tail in (".nwk", ".tre", ".tree", ".newick", ".nex"):
            return head
        else:
            return fname

    if (not format):
        if (type(data) in StringTypes) and os.path.isfile(data):
            s = data.lower()
            for tail in ".nex", ".nexus", ".tre":
                if s.endswith(tail):
                    format="nexus"
                    break

    if (not format):
        format = "newick"

    if format == "newick":
        if type(data) in StringTypes:
            if os.path.isfile(data):
                treename = strip(data)
                return newick.parse(file(data), treename=treename,
                                    ttable=ttable)
            else:
                return newick.parse(data, ttable=ttable)

        elif (hasattr(data, "tell") and hasattr(data, "read")):
            treename = strip(getattr(data, "name", None))
            return newick.parse(data, treename=treename, ttable=ttable)
    elif format == "nexus-dendropy":
        import dendropy
        if type(data) in StringTypes:
            if os.path.isfile(data):
                treename = strip(data)
                return newick.parse(
                    str(dendropy.Tree.get_from_path(data, "nexus")),
                    treename=treename
                    )
            else:
                return newick.parse(
                    str(dendropy.Tree.get_from_string(data, "nexus"))
                    )

        elif (hasattr(data, "tell") and hasattr(data, "read")):
            treename = strip(getattr(data, "name", None))
            return newick.parse(
                str(dendropy.Tree.get_from_stream(data, "nexus")),
                treename=treename
                )
        else:
            pass

    elif format == "nexus":
        if type(data) in StringTypes:
            if os.path.isfile(data):
                with open(data) as infile:
                    rec = newick.nexus_iter(infile).next()
                    if rec: return rec.parse()
            else:
                rec = newick.nexus_iter(StringIO(data)).next()
                if rec: return rec.parse()
        else:
            rec = newick.nexus_iter(data).next()
            if rec: return rec.parse()
    else:
        # implement other tree formats here (nexus, nexml etc.)
        raise IOError, "format '%s' not implemented yet" % format

    raise IOError, "unable to read tree from '%s'" % data

def readmany(data, format="newick"):
    """Iterate over trees from a source."""
    if type(data) in types.StringTypes:
        if os.path.isfile(data):
            data = open(data)
        else:
            data = StringIO(data)

    if format == "newick":
        for line in data:
            yield newick.parse(line)
    elif format == "nexus":
        for rec in newick.nexus_iter(data):
            yield rec.parse()
    else:
        raise Exception, "format '%s' not recognized" % format
    data.close()

## def randomly_resolve(n):
##     assert len(n.children)>2

## def leaf_mrcas(root):
##     from itertools import product, izip, tee
##     from collections import OrderedDict
##     from numpy import empty
##     mrca = OrderedDict()
##     def pairwise(iterable, tee=tee, izip=izip):
##         a, b = tee(iterable)
##         next(b, None)
##         return izip(a, b)
##     def f(n):
##         if n.isleaf:
##             od = OrderedDict(); od[n] = n.length
##             return od
##         d = [ f(c) for c in n.children ]
##         for i, j in pairwise(xrange(len(d))):
##             di = d[i]; dj =d[j]
##             for ni, niv in di.iteritems():
##                 for nj, njv in dj.iteritems():
##                     mrca[(ni,nj)] = n
##             d[j].update(di)
##         return d[j]
##     f(root)
##     return mrca

def C(leaves, internals):
    from scipy.sparse import lil_matrix
    m = lil_matrix((len(internals), len(leaves)))
    for lf in leaves:
        v = lf.length if lf.length is not None else 1
        for n in lf.rootpath():
            m[n.ii,lf.li] = v
            v += n.length if n.length is not None else 1
    return m.tocsc()

"""
Functions to get trees and character data from treebase
"""

from urllib2 import urlopen
from lxml import etree
from collections import defaultdict
from storage import Storage
import sys, re

# "http://purl.org/phylo/treebase/phylows/study/TB2:S11152"

TREEBASE_WEBSERVICE = "http://purl.org/phylo/treebase/phylows"
NEXML_NAMESPACE = "http://www.nexml.org/2009"
NEXML = "{%s}" % NEXML_NAMESPACE
UNIPROT = "http://purl.uniprot.org/taxonomy/"
NAMEBANK = ("http://www.ubio.org/authority/metadata.php?"
            "lsid=urn:lsid:ubio.org:namebank:")

ROW_SEGMENTS = ("http://treebase.org/treebase-web/search/study/"
                "rowSegmentsTSV.html?matrixid=")

META_DATATYPE = {
    "xsd:long": int,
    "xsd:integer": int,
    "xsd:string": str
    }

AMBIG_RE = re.compile(r'([{][a-zA-Z]+[}])')

def fetch_study(study_id, format="nexml"):
    """
    Get a study from treebase in one of various formats

    Args:
        study_id (str): The id of the study
        format (str):  One of ["rdf", "html", "nexml", "nexus"]
    Returns:
        Str representing a nexus file (if format = "nexus")

        OR

        An lxml etree object
    """
    try: study_id = "S%s" % int(study_id)
    except ValueError: pass

    # format is one of ["rdf", "html", "nexml", "nexus"]
    url = "%s/study/TB2:%s?format=%s" % (TREEBASE_WEBSERVICE, study_id, format)
    if format=="nexus":
        return urlopen(url).read()
    else:
        return etree.parse(url)

def parse_chars(e, otus):
    v = []
    for chars in e.findall(NEXML+"characters"):
        c = Storage(chars.attrib)
        c.states = parse_states(chars)
        c.meta = Storage()
        for meta in chars.findall(NEXML+"meta"):
            a = meta.attrib
            if a.get("content"):
                value = META_DATATYPE[a["datatype"]](a["content"])
                c.meta[a["property"]] = value
        c.matrices = []
        for matrix in chars.findall(NEXML+"matrix"):
            m = Storage()
            m.rows = []
            for row in matrix.findall(NEXML+"row"):
                r = Storage(row.attrib)
                r.otu = otus[r.otu]
                s = row.findall(NEXML+"seq")[0].text
                substrs = []
                for ss in AMBIG_RE.split(s):
                    if ss.startswith("{"):
                        key = frozenset(ss[1:-1])
                        val = c.states.states2symb.get(key)
                        if key and not val:
                            sys.stderr.write("missing ambig symbol for %s\n" %
                                             "".join(sorted(key)))
                        ss = val or "?"
                    substrs.append(ss)
                s = "".join(substrs)
                r.seq = s
                m.rows.append(r)
            c.matrices.append(m)
        v.append(c)
    return v

def parse_trees(e, otus):
    """
    Get trees from an etree object

    Args:
        e: A nexml document parsed by etree
        otus: OTUs returned by parse_otus
    Returns:
        list: A list of ivy Storage objects each
          containing every node of a tree.
    """
    from tree import Node
    v = []
    for tb in e.findall(NEXML+"trees"):
        for te in tb.findall(NEXML+"tree"):
            t = Storage()
            t.attrib = Storage(te.attrib)
            t.nodes = {}
            for n in te.findall(NEXML+"node"):
                node = Node()
                if n.attrib.get("otu"):
                    node.isleaf = True
                    node.otu = otus[n.attrib["otu"]]
                    node.label = node.otu.label
                t.nodes[n.attrib["id"]] = node
            for edge in te.findall(NEXML+"edge"):
                d = edge.attrib
                n = t.nodes[d["target"]]
                p = t.nodes[d["source"]]
                length = d.get("length")
                if length:
                    n.length = float(length)
                p.add_child(n)
            r = [ n for n in t.nodes.values() if not n.parent ]
            assert len(r)==1
            r = r[0]
            r.isroot = True
            for i, n in enumerate(r): n.id = i+1
            t.root = r
            v.append(t)
    return v

def parse_otus(e):
    """
    Get OTUs from an etree object

    Args:
        e: A nexml document parsed by etree
    Returns:
        dict: A dict mapping keys to OTUs contained in ivy Storage objects
    """
    v = {}
    for otus in e.findall(NEXML+"otus"):
        for x in otus.findall(NEXML+"otu"):
            otu = Storage()
            otu.id = x.attrib["id"]
            otu.label = x.attrib["label"]
            for meta in x.iterchildren():
                d = meta.attrib
                p = d.get("property")
                if p and p == "tb:identifier.taxon":
                    otu.tb_taxid = d["content"]
                elif p and p == "tb:identifier.taxonVariant":
                    otu.tb_taxid_variant = d["content"]
                h = d.get("href")
                if h and h.startswith(NAMEBANK):
                    otu.namebank_id = int(h.replace(NAMEBANK, ""))
                elif h and h.startswith(UNIPROT):
                    otu.ncbi_taxid = int(h.replace(UNIPROT, ""))
            v[otu.id] = otu
    return v

def parse_nexml(doc):
    """
    Parse an etree ElementTree

    Args:
        doc: An etree ElementTree or a file that can be parsed into
          an etree ElementTree with etree.parse
    Returns:
        An ivy Storage object containing all the information from the
          nexml file: Characters, metadata, OTUs, and trees.
    """
    if not isinstance(doc, (etree._ElementTree, etree._Element)):
        doc = etree.parse(doc)
    meta = {}
    for child in doc.findall(NEXML+"meta"):
        if "content" in child.attrib:
            d = child.attrib
            key = d["property"]
            val = META_DATATYPE[d["datatype"]](d["content"])
            if (key in meta) and val:
                if isinstance(meta[key], list):
                    meta[key].append(val)
                else:
                    meta[key] = [meta[key], val]
            else:
                meta[key] = val

    otus = parse_otus(doc)

    return Storage(meta = meta,
                   otus = otus,
                   chars = parse_chars(doc, otus),
                   trees = parse_trees(doc, otus))

def parse_states(e):
    """e is a characters element"""
    f = e.findall(NEXML+"format")[0]
    sts = f.findall(NEXML+"states")[0]
    states2symb = {}
    symb2states = {}
    id2symb = {}
    for child in sts.iterchildren():
        t = child.tag
        if t == NEXML+"state":
            k = child.attrib["id"]
            v = child.attrib["symbol"]
            id2symb[k] = v
            states2symb[v] = v
            symb2states[v] = v
        elif t == NEXML+"uncertain_state_set":
            k = child.attrib["id"]
            v = child.attrib["symbol"]
            id2symb[k] = v
            memberstates = []
            for memb in child.findall(NEXML+"member"):
                sid = memb.attrib["state"]
                symb = id2symb[sid]
                for x in symb2states[symb]: memberstates.append(x)
            memberstates = frozenset(memberstates)
            symb2states[v] = memberstates
            states2symb[memberstates] = v
    return Storage(states2symb=states2symb,
                   symb2states=symb2states,
                   id2symb=id2symb)

def parse_charsets(study_id):
    from cStringIO import StringIO
    nx = StringIO(fetch_study(study_id, 'nexus'))
    d = {}
    for line in nx.readlines():
        if line.strip().startswith("CHARSET "):
            v = line.strip().split()
            label = v[1]
            first, last = map(int, line.split()[-1][:-1].split("-"))
            d[label] = (first-1, last-1)
    return d

if __name__ == "__main__":
    import sys
    from pprint import pprint
    e = fetch_study('S11152', 'nexus')
    #print e
    #e.write(sys.stdout, pretty_print=True)

    ## e = etree.parse('/tmp/tmp.xml')
    ## x = parse_nexml(e)
    ## pprint(x)

"""
interactive viewers for trees, etc. using matplotlib
"""
from collections import defaultdict
from ..storage import Storage
from .. import align, sequtil
import matplotlib, numpy, types
import matplotlib.pyplot as pyplot
from matplotlib.figure import SubplotParams, Figure
from matplotlib.axes import Axes, subplot_class_factory
from matplotlib.patches import Rectangle
from matplotlib.collections import PatchCollection
from matplotlib.widgets import RectangleSelector
from matplotlib.transforms import Bbox, offset_copy, IdentityTransform
from matplotlib import colors as mpl_colors
from matplotlib.ticker import MaxNLocator, FuncFormatter, NullLocator
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from Bio.Align import MultipleSeqAlignment

matplotlib.rcParams['path.simplify'] = False

class UpdatingRect(Rectangle):
    def __call__(self, p):
        self.set_bounds(*p.viewLim.bounds)
        p.figure.canvas.draw_idle()

class AlignmentFigure:
    def __init__(self, aln, name=None, div=0.25, overview=True):
        if isinstance(aln, MultipleSeqAlignment):
            self.aln = aln
        else:
            self.aln = align.read(aln)
        self.name = name
        self.div_value = div
        pars = SubplotParams(
            left=0.2, right=1, bottom=0.05, top=1, wspace=0.01
            )
        fig = pyplot.figure(subplotpars=pars, facecolor="white")
        self.figure = fig
        self.initialize_subplots(overview)
        self.show()
        self.connect_events()

    def initialize_subplots(self, overview=False):
        ## p = AlignmentPlot(self.figure, 212, aln=self.aln)
        p = AlignmentPlot(self.figure, 111, aln=self.aln, app=self)
        self.detail = self.figure.add_subplot(p)
        self.detail.plot_aln()
        if overview:
            self.overview = inset_axes(
                self.detail, width="30%", height="20%", loc=1
                )
            self.overview.xaxis.set_major_locator(NullLocator())
            self.overview.yaxis.set_major_locator(NullLocator())
            self.overview.imshow(
                self.detail.array, interpolation='nearest', aspect='auto',
                origin='lower'
                )
            rect = UpdatingRect(
                [0,0], 0, 0, facecolor='black', edgecolor='cyan', alpha=0.5
                )
            self.overview.zoomrect = rect
            rect.target = self.detail
            self.detail.callbacks.connect('xlim_changed', rect)
            self.detail.callbacks.connect('ylim_changed', rect)
            self.overview.add_patch(rect)
            rect(self.overview)

        else:
            self.overview = None

    def show(self):
        self.figure.show()

    def connect_events(self):
        mpl_connect = self.figure.canvas.mpl_connect
        mpl_connect("button_press_event", self.onclick)
        mpl_connect("button_release_event", self.onbuttonrelease)
        mpl_connect("scroll_event", self.onscroll)
        mpl_connect("pick_event", self.onpick)
        mpl_connect("motion_notify_event", self.ondrag)
        mpl_connect("key_press_event", self.onkeypress)
        mpl_connect("axes_enter_event", self.axes_enter)
        mpl_connect("axes_leave_event", self.axes_leave)

    @staticmethod
    def axes_enter(e):
        ax = e.inaxes
        ax._active = True

    @staticmethod
    def axes_leave(e):
        ax = e.inaxes
        ax._active = False

    @staticmethod
    def onselect(estart, estop):
        b = estart.button
        ## print b, estart.key

    @staticmethod
    def onkeypress(e):
        ax = e.inaxes
        k = e.key
        if ax and k:
            if k == 't':
                ax.home()
            elif k == "down":
                ax.scroll(0, -0.1)
            elif k == "up":
                ax.scroll(0, 0.1)
            elif k == "left":
                ax.scroll(-0.1, 0)
            elif k == "right":
                ax.scroll(0.1, 0)
            elif k in '=+':
                ax.zoom(0.1,0.1)
            elif k == '-':
                ax.zoom(-0.1,-0.1)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def ondrag(e):
        ax = e.inaxes
        button = e.button
        if ax and button == 2:
            if not ax.pan_start:
                ax.pan_start = (e.xdata, e.ydata)
                return
            x, y = ax.pan_start
            xdelta = x - e.xdata
            ydelta = y - e.ydata
            x0, x1 = ax.get_xlim()
            xspan = x1-x0
            y0, y1 = ax.get_ylim()
            yspan = y1 - y0
            midx = (x1+x0)*0.5
            midy = (y1+y0)*0.5
            ax.set_xlim(midx+xdelta-xspan*0.5, midx+xdelta+xspan*0.5)
            ax.set_ylim(midy+ydelta-yspan*0.5, midy+ydelta+yspan*0.5)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onbuttonrelease(e):
        ax = e.inaxes
        button = e.button
        if button == 2:
            ## print "pan end"
            ax.pan_start = None
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onpick(e):
        ax = e.mouseevent.inaxes
        if ax:
            ax.picked(e)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onscroll(e):
        ax = e.inaxes
        if ax:
            b = e.button
            ## print b
            k = e.key
            if k == None and b =="up":
                ax.zoom(0.1,0.1)
            elif k == None and b =="down":
                ax.zoom(-0.1,-0.1)
            elif k == "shift" and b == "up":
                ax.zoom_cxy(0.1, 0, e.xdata, e.ydata)
            elif k == "shift" and b == "down":
                ax.zoom_cxy(-0.1, 0, e.xdata, e.ydata)
            elif k == "control" and b == "up":
                ax.zoom_cxy(0, 0.1, e.xdata, e.ydata)
            elif k == "control" and b == "down":
                ax.zoom_cxy(0, -0.1, e.xdata, e.ydata)
            elif k == "d" and b == "up":
                ax.scroll(0, 0.1)
            elif (k == "d" and b == "down"):
                ax.scroll(0, -0.1)
            elif k == "c" and b == "up":
                ax.scroll(-0.1, 0)
            elif k == "c" and b == "down":
                ax.scroll(0.1, 0)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onclick(e):
        ax = e.inaxes
        if (ax and e.button==1 and hasattr(ax, "zoomrect") and ax.zoomrect):
            # overview clicked; reposition zoomrect
            r = ax.zoomrect
            x = e.xdata
            y = e.ydata
            arr = ax.transData.inverted().transform(r.get_extents())
            xoff = (arr[1][0]-arr[0][0])*0.5
            yoff = (arr[1][1]-arr[0][1])*0.5
            r.target.set_xlim(x-xoff,x+xoff)
            r.target.set_ylim(y-yoff,y+yoff)
            r(r.target)
            ax.figure.canvas.draw_idle()

        elif ax and e.button==2:
            ## print "pan start", (e.xdata, e.ydata)
            ax.pan_start = (e.xdata, e.ydata)
            ax.figure.canvas.draw_idle()

        elif ax and hasattr(ax, "aln") and ax.aln:
            x = int(e.xdata+0.5); y = int(e.ydata+0.5)
            aln = ax.aln
            if (0 <= x <= ax.nchar) and (0 <= y <= ax.ntax):
                seq = aln[y]; char = seq[x]
                if char not in '-?':
                    v = sequtil.gapidx(seq)
                    i = (v[1]==x).nonzero()[0][0]
                    print ("%s: row %s, site %s: '%s', seqpos %s"
                           % (seq.id, y, x, char, i))
                else:
                    print "%s: row %s, site %s: '%s'" % (seq.id, y, x, char)

    def zoom(self, factor=0.1):
        "Zoom both axes by *factor* (relative display size)."
        self.detail.zoom(factor, factor)
        self.figure.canvas.draw_idle()

    def __get_selection(self):
        return self.detail.extract_selected()
    selected = property(__get_selection)

class Alignment(Axes):
    """
    matplotlib.axes.Axes subclass for rendering sequence alignments.
    """
    def __init__(self, fig, rect, *args, **kwargs):
        self.aln = kwargs.pop("aln")
        nrows = len(self.aln)
        ncols = self.aln.get_alignment_length()
        self.alnidx = numpy.arange(ncols)
        self.app = kwargs.pop("app", None)
        self.showy = kwargs.pop('showy', True)
        Axes.__init__(self, fig, rect, *args, **kwargs)
        rgb = mpl_colors.colorConverter.to_rgb
        gray = rgb('gray')
        d = defaultdict(lambda:gray)
        d["A"] = rgb("red")
        d["a"] = rgb("red")
        d["C"] = rgb("blue")
        d["c"] = rgb("blue")
        d["G"] = rgb("green")
        d["g"] = rgb("green")
        d["T"] = rgb("yellow")
        d["t"] = rgb("yellow")
        self.cmap = d
        self.selector = RectangleSelector(
            self, self.select_rectangle, useblit=True
            )
        def f(e):
            if e.button != 1: return True
            else: return RectangleSelector.ignore(self.selector, e)
        self.selector.ignore = f
        self.selected_rectangle = Rectangle(
            [0,0],0,0, facecolor='white', edgecolor='cyan', alpha=0.3
            )
        self.add_patch(self.selected_rectangle)
        self.highlight_find_collection = None

    def plot_aln(self):
        cmap = self.cmap
        self.ntax = len(self.aln); self.nchar = self.aln.get_alignment_length()
        a = numpy.array([ [ cmap[base] for base in x.seq ]
                          for x in self.aln ])
        self.array = a
        self.imshow(a, interpolation='nearest', aspect='auto', origin='lower')
        y = [ i+0.5 for i in xrange(self.ntax) ]
        labels = [ x.id for x in self.aln ]
        ## locator.bin_boundaries(1,ntax)
        ## locator.view_limits(1,ntax)
        if self.showy:
            locator = MaxNLocator(nbins=50, integer=True)
            self.yaxis.set_major_locator(locator)
            def fmt(x, pos=None):
                if x<0: return ""
                try: return labels[int(round(x))]
                except: pass
                return ""
            self.yaxis.set_major_formatter(FuncFormatter(fmt))
        else:
            self.yaxis.set_major_locator(NullLocator())

        return self

    def select_rectangle(self, e0, e1):
        x0, x1 = map(int, sorted((e0.xdata+0.5, e1.xdata+0.5)))
        y0, y1 = map(int, sorted((e0.ydata+0.5, e1.ydata+0.5)))
        self.selected_chars = (x0, x1)
        self.selected_taxa = (y0, y1)
        self.selected_rectangle.set_bounds(x0-0.5,y0-0.5,x1-x0+1,y1-y0+1)
        self.app.figure.canvas.draw_idle()

    def highlight_find(self, substr):
        if not substr:
            if self.highlight_find_collection:
                self.highlight_find_collection.remove()
                self.highlight_find_collection = None
            return

        N = len(substr)
        v = []
        for y, x in align.find(self.aln, substr):
            r = Rectangle(
                [x-0.5,y-0.5], N, 1,
                facecolor='cyan', edgecolor='cyan', alpha=0.7
                )
            v.append(r)
        if self.highlight_find_collection:
            self.highlight_find_collection.remove()
        c = PatchCollection(v, True)
        self.highlight_find_collection = self.add_collection(c)
        self.app.figure.canvas.draw_idle()

    def extract_selected(self):
        r0, r1 = self.selected_taxa
        c0, c1 = self.selected_chars
        return self.aln[r0:r1+1,c0:c1+1]

    def zoom_cxy(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view, with a fixed data point (cx, cy)
        """
        transform = self.transData.inverted().transform
        xlim = self.get_xlim(); xmid = sum(xlim)*0.5
        ylim = self.get_ylim(); ymid = sum(ylim)*0.5
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = xmid-x0; deltay = ymid-y0
        cx = cx or xmid; cy = cy or ymid
        xoff = (cx-xmid)*x
        self.set_xlim(xmid-deltax+xoff, xmid+deltax+xoff)
        yoff = (cy-ymid)*y
        self.set_ylim(ymid-deltay+yoff, ymid+deltay+yoff)

    def zoom(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view.
        """
        # get the function to convert display coordinates to data
        # coordinates
        transform = self.transData.inverted().transform
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = x0 - xlim[0]; deltay = y0 - ylim[0]
        self.set_xlim(xlim[0]+deltax, xlim[1]-deltax)
        self.set_ylim(ylim[0]+deltay, ylim[1]-deltay)

    def center_y(self, y):
        ymin, ymax = self.get_ylim()
        yoff = (ymax - ymin) * 0.5
        self.set_ylim(y-yoff, y+yoff)

    def center_x(self, x, offset=0.3):
        xmin, xmax = self.get_xlim()
        xspan = xmax - xmin
        xoff = xspan*0.5 + xspan*offset
        self.set_xlim(x-xoff, x+xoff)

    def scroll(self, x, y):
        x0, x1 = self.get_xlim()
        y0, y1 = self.get_ylim()
        xd = (x1-x0)*x
        yd = (y1-y0)*y
        self.set_xlim(x0+xd, x1+xd)
        self.set_ylim(y0+yd, y1+yd)

    def home(self):
        self.set_xlim(0, self.nchar)
        self.set_ylim(self.ntax, 0)

AlignmentPlot = subplot_class_factory(Alignment)

import os, matplotlib
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
import tree
from axes_utils import adjust_limits
import tempfile

## class TreeFigure:
##     def __init__(self):
##         pass

matplotlib.rcParams["xtick.direction"] = "out"

class TreeFigure:
    def __init__(self, root, relwidth=0.5, leafpad=1.5, name=None,
                 support=70.0, scaled=True, mark_named=True,
                 leaf_fontsize=10, branch_fontsize=10,
                 branch_width=1, branch_color="black",
                 highlight_support=True,
                 branchlabels=True, leaflabels=True, decorators=[],
                 xoff=0, yoff=0,
                 xlim=None, ylim=None,
                 height=None, width=None):
        self.root = root
        self.relwidth = relwidth
        self.leafpad = leafpad
        self.name = name
        self.support = support
        self.scaled = scaled
        self.mark_named = mark_named
        self.leaf_fontsize = leaf_fontsize
        self.branch_fontsize = branch_fontsize
        self.branch_width = branch_width
        self.branch_color = branch_color
        self.highlight_support = highlight_support
        self.branchlabels = branchlabels
        self.leaflabels = leaflabels
        self.decorators = decorators
        self.xoff = xoff
        self.yoff = yoff

        nleaves = len(root.leaves())
        self.dpi = 72.0
        h = height or (nleaves*self.leaf_fontsize*self.leafpad)/self.dpi
        self.height = h
        self.width = width or self.height*self.relwidth
        ## p = min(self.width, self.height)*0.1
        ## self.height += p
        ## self.width += p
        self.figure = Figure(figsize=(self.width, self.height), dpi=self.dpi)
        self.canvas = FigureCanvas(self.figure)
        self.axes = self.figure.add_axes(
            tree.TreePlot(self.figure, 1,1,1,
                          support=self.support,
                          scaled=self.scaled,
                          mark_named=self.mark_named,
                          leaf_fontsize=self.leaf_fontsize,
                          branch_fontsize=self.branch_fontsize,
                          branch_width=self.branch_width,
                          branch_color=self.branch_color,
                          highlight_support=self.highlight_support,
                          branchlabels=self.branchlabels,
                          leaflabels=self.leaflabels,
                          interactive=False,
                          decorators=self.decorators,
                          xoff=self.xoff, yoff=self.yoff,
                          name=self.name).plot_tree(self.root)
            )
        self.axes.spines["top"].set_visible(False)
        self.axes.spines["left"].set_visible(False)
        self.axes.spines["right"].set_visible(False)
        self.axes.spines["bottom"].set_smart_bounds(True)
        self.axes.xaxis.set_ticks_position("bottom")

        for v in self.axes.node2label.values():
            v.set_visible(True)

        ## for k, v in self.decorators:
        ##     func, args, kwargs = v
        ##     func(self.axes, *args, **kwargs)

        self.canvas.draw()
        ## self.axes.home()
        ## adjust_limits(self.axes)
        self.axes.set_position([0.05,0.05,0.95,0.95])

    @property
    def detail(self):
        return self.axes

    def savefig(self, fname):
        root, ext = os.path.splitext(fname)
        buf = tempfile.TemporaryFile()
        for i in range(3):
            self.figure.savefig(buf, format=ext[1:].lower())
            self.home()
            buf.seek(0)
        buf.close()
        self.figure.savefig(fname)

    def set_relative_width(self, relwidth):
        w, h = self.figure.get_size_inches()
        self.figure.set_figwidth(h*relwidth)

    def autoheight(self):
        "adjust figure height to show all leaf labels"
        nleaves = len(self.root.leaves())
        h = (nleaves*self.leaf_fontsize*self.leafpad)/self.dpi
        self.height = h
        self.figure.set_size_inches(self.width, self.height)
        self.axes.set_ylim(-2, nleaves+2)

    def home(self):
        self.axes.home()

"""
Convenience functions for drawing shapes on TreePlots.
"""
try:
    import Image
except ImportError:
    from PIL import Image
from numpy import pi
from matplotlib.collections import RegularPolyCollection, CircleCollection
from matplotlib.transforms import offset_copy
from matplotlib.patches import Rectangle, Wedge, Circle, PathPatch
from matplotlib.offsetbox import DrawingArea
from itertools import izip_longest
from matplotlib.axes import Axes
from numpy import array
from matplotlib.path import Path

try:
    from matplotlib.offsetbox import OffsetImage, AnnotationBbox
except ImportError:
    pass
from ..tree import Node
import colors as _colors

def _xy(plot, p):
    if isinstance(p, Node):
        c = plot.n2c[p]
        p = (c.x, c.y)
    elif isinstance(p, (list, tuple)):
        p = [ _xy(plot, x) for x in p ]
    else:
        pass
    return p

def image(plot, p, imgfile,
          maxdim=100, border=0,
          xoff=4, yoff=4,
          halign=0.0, valign=0.5,
          xycoords='data',
          boxcoords=('offset points')):
    """
    Add images to plot

    Args:
        plot (Tree): A Tree plot instance
        p (Node): A node object
        imgfile (str): A path to an image
        maxdim (float): Maximum dimension of image. Optional,
          defaults to 100.
        border: RR: What does border do? -CZ
        xoff, yoff (float): X and Y offset. Optional, defaults to 4
        halign, valign (float): Horizontal and vertical alignment within
          box. Optional, defaults to 0.0 and 0.5, respectively.

    """
    if xycoords == "label":
        xycoords = plot.node2label[p]
        x, y = (1, 0.5)
    else:
        x, y = _xy(plot, p)
    img = Image.open(imgfile)
    if max(img.size) > maxdim:
        img.thumbnail((maxdim, maxdim))
    imgbox = OffsetImage(img)
    abox = AnnotationBbox(imgbox, (x, y),
                          xybox=(xoff, yoff),
                          xycoords=xycoords,
                          box_alignment=(halign,valign),
                          pad=0.0,
                          boxcoords=boxcoords)
    plot.add_artist(abox)
    plot.figure.canvas.draw_idle()

def images(plot, p, imgfiles,
          maxdim=100, border=0,
          xoff=4, yoff=4,
          halign=0.0, valign=0.5,
          xycoords='data', boxcoords=('offset points')):
    """
    Add many images to plot at once

    Args:
        Plot (Tree): A Tree plot instance
        p (list): A list of node objects
        imgfile (list): A list of strs containing paths to image files.
          Must be the same length as p.
        maxdim (float): Maximum dimension of image. Optional,
          defaults to 100.
        border: RR: What does border do? -CZ
        xoff, yoff (float): X and Y offset. Optional, defaults to 4
        halign, valign (float): Horizontal and vertical alignment within
          box. Optional, defaults to 0.0 and 0.5, respectively.

    """
    for x, f in zip(p, imgfiles):
        image(plot, x, f, maxdim, border, xoff, yoff, halign, valign,
              xycoords, boxcoords)

def pie(plot, p, values, colors=None, size=16, norm=True,
        xoff=0, yoff=0,
        halign=0.5, valign=0.5,
        xycoords='data', boxcoords=('offset points')):
    """
    Draw a pie chart

    Args:
    plot (Tree): A Tree plot instance
    p (Node): A Node object
    values (list): A list of floats.
    colors (list): A list of strings to pull colors from. Optional.
    size (float): Diameter of the pie chart
    norm (bool): Whether or not to normalize the values so they
      add up to 360
    xoff, yoff (float): X and Y offset. Optional, defaults to 0
    halign, valign (float): Horizontal and vertical alignment within
      box. Optional, defaults to 0.5

    """
    x, y = _xy(plot, p)
    da = DrawingArea(size, size); r = size*0.5; center = (r,r)
    x0 = 0
    S = 360.0
    if norm: S = 360.0/sum(values)
    if not colors:
        c = _colors.tango()
        colors = [ c.next() for v in values ]
    for i, v in enumerate(values):
        theta = v*S
        if v: da.add_artist(Wedge(center, r, x0, x0+theta,
                                  fc=colors[i], ec='none'))
        x0 += theta
    box = AnnotationBbox(da, (x,y), pad=0, frameon=False,
                         xybox=(xoff, yoff),
                         xycoords=xycoords,
                         box_alignment=(halign,valign),
                         boxcoords=boxcoords)
    plot.add_artist(box)
    plot.figure.canvas.draw_idle()
    return box

def hbar(plot, p, values, colors=None, height=16,
         xoff=0, yoff=0,
         halign=1, valign=0.5,
         xycoords='data', boxcoords=('offset points')):
    x, y = _xy(plot, p)
    h = height; w = sum(values) * height#; yoff=h*0.5
    da = DrawingArea(w, h)
    x0 = -sum(values)
    if not colors:
        c = _colors.tango()
        colors = [ c.next() for v in values ]
    for i, v in enumerate(values):
        if v: da.add_artist(Rectangle((x0,0), v*h, h, fc=colors[i], ec='none'))
        x0 += v*h
    box = AnnotationBbox(da, (x,y), pad=0, frameon=False,
                         xybox=(xoff, yoff),
                         xycoords=xycoords,
                         box_alignment=(halign,valign),
                         boxcoords=boxcoords)
    plot.add_artist(box)
    plot.figure.canvas.draw_idle()

def hbars(plot, p, values, colors=None, height=16,
          xoff=0, yoff=0,
          halign=1, valign=0.5,
          xycoords='data', boxcoords=('offset points')):
    for x, v in zip(p, values):
        hbar(plot, x, v, colors, height, xoff, yoff, halign, valign,
             xycoords, boxcoords)

def squares(plot, p, colors='r', size=15, xoff=0, yoff=0, alpha=1.0,
            zorder=1000):
    """
    Draw a square at given node

    Args:
        plot (Tree): A Tree plot instance
        p: A node or list of nodes
        colors: Str or list of strs. Colors of squares to be drawn.
          Optional, defaults to 'r' (red)
        size (float): Size of the squares. Optional, defaults to 15
        xoff, yoff (float): Offset for x and y dimensions. Optional,
          defaults to 0.
        alpha (float): between 0 and 1. Alpha transparency of squares.
          Optional, defaults to 1 (fully opaque)
        zorder (int): The drawing order. Higher numbers appear on top
          of lower numbers. Optional, defaults to 1000.

    """
    points = _xy(plot, p)
    trans = offset_copy(
        plot.transData, fig=plot.figure, x=xoff, y=yoff, units='points')

    col = RegularPolyCollection(
        numsides=4, rotation=pi*0.25, sizes=(size*size,),
        offsets=points, facecolors=colors, transOffset=trans,
        edgecolors='none', alpha=alpha, zorder=zorder
        )

    plot.add_collection(col)
    plot.figure.canvas.draw_idle()

def tipsquares(plot, p, colors='r', size=15, pad=2, edgepad=10):
    """
    RR: Bug with this function. If you attempt to call it with a list as an
    argument for p, it will not only not work (expected) but it will also
    make it so that you can't interact with the tree figure (gives errors when
    you try to add symbols, select nodes, etc.) -CZ

    Add square after tip label, anchored to the side of the plot

    Args:
        plot (Tree): A Tree plot instance.
        p (Node): A Node object (Should be a leaf node).
        colors (str): olor of drawn square. Optional, defaults to 'r' (red)
        size (float): Size of square. Optional, defaults to 15
        pad: RR: I am unsure what this does. Does not seem to have visible
          effect when I change it. -CZ
        edgepad (float): Padding from square to edge of plot. Optional,
          defaults to 10.

    """
    x, y = _xy(plot, p) # p is a single node or point in data coordinates
    n = len(colors)
    da = DrawingArea(size*n+pad*(n-1), size, 0, 0)
    sx = 0
    for c in colors:
        sq = Rectangle((sx,0), size, size, color=c)
        da.add_artist(sq)
        sx += size+pad
    box = AnnotationBbox(da, (x, y), xybox=(-edgepad,y),
                         frameon=False,
                         pad=0.0,
                         xycoords='data',
                         box_alignment=(1, 0.5),
                         boxcoords=('axes points','data'))
    plot.add_artist(box)
    plot.figure.canvas.draw_idle()

def circles(plot, p, colors='g', size=15, xoff=0, yoff=0):
    """
    Draw circles on plot

    Args:
        plot (Tree): A Tree plot instance
        p: A node object or list of Node objects
        colors: Str or list of strs. Colors of the circles. Optional,
          defaults to 'g' (green)
        size (float): Size of the circles. Optional, defaults to 15
        xoff, yoff (float): X and Y offset. Optional, defaults to 0.

    """
    points = _xy(plot, p)
    trans = offset_copy(
        plot.transData, fig=plot.figure, x=xoff, y=yoff, units='points'
        )

    col = CircleCollection(
        sizes=(pi*size*size*0.25,),
        offsets=points, facecolors=colors, transOffset=trans,
        edgecolors='none'
        )

    plot.add_collection(col)
    plot.figure.canvas.draw_idle()
    return col

def legend(plot, colors, labels, shape='rectangle', loc='upper left', **kwargs):
    """
    RR: the MPL legend function has changed since this function has been
    written. This function currently does not work. -CZ
    """
    if shape == 'circle':
        shapes = [ Circle((0.5,0.5), radius=1, fc=c) for c in colors ]
        #shapes = [ CircleCollection([10],facecolors=[c]) for c in colors ]
    else:
        shapes = [ Rectangle((0,0),1,1,fc=c,ec='none') for c in colors ]

    return Axes.legend(plot, shapes, labels, loc=loc, **kwargs)

def leafspace_triangles(plot, color='black', rca=0.5):
    """
    RR: Using this function on the primates tree (straight from the newick file)
    gives error: 'Node' object has no attribute 'leafspace'. How do you give
    nodes the leafspace attribute? -CZ
    rca = relative crown age
    """
    leaves = plot.root.leaves()
    leafspace = [ float(x.leafspace) for x in leaves ]
    #leafspace = array(raw_leafspace)/(sum(raw_leafspace)/float(len(leaves)))
    pv = []
    for i, n in enumerate(leaves):
        if leafspace[i] > 0:
            p = plot.n2c[n]
            pp = plot.n2c[n.parent]
            spc = leafspace[i]
            yoff = spc/2.0
            x0 = pp.x + (p.x - pp.x)*rca
            verts = [(x0, p.y),
                     (p.x, p.y-yoff),
                     (p.x, p.y+yoff),
                     (x0, p.y)]
            codes = [Path.MOVETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY]
            path = Path(verts, codes)
            patch = PathPatch(path, fc=color, lw=0)
            pv.append(plot.add_patch(patch))
    return pv

def text(plot, x, y, s, color='black', xoff=0, yoff=0, valign='center',
         halign='left', fontsize=10):
    """
    Add text to the plot.

    Args:
        plot (Tree): A Tree plot instance
        x, y (float): x and y coordinates to place the text
        s (str): The text to write
        color (str): The color of the text. Optional, defaults to "black"
        xoff, yoff (float): x and y offset
        valign (str): Vertical alignment. Can be: 'center', 'top',
          'bottom', or 'baseline'. Defaults to 'center'.
        halign (str): Horizontal alignment. Can be: 'center', 'right',
          or 'left'. Defaults to 'left'
        fontsize (float): Font size. Optional, defaults to 10

    """
    txt = plot.annotate(
        s, xy=(x, y),
        xytext=(xoff, yoff),
        textcoords="offset points",
        verticalalignment=valign,
        horizontalalignment=halign,
        fontsize=fontsize,
        clip_on=True,
        picker=True
    )
    txt.set_visible(True)
    return txt

"""
interactive viewers for trees, etc. using matplotlib
"""
import sys, time, bisect, math, types, os, operator
from collections import defaultdict
from itertools import chain
from pprint import pprint
from .. import tree, bipart
from ..layout import cartesian
from ..storage import Storage
from .. import pyperclip as clipboard
#from ..nodecache import NodeCache
import matplotlib, numpy
import matplotlib.pyplot as pyplot
from matplotlib.figure import SubplotParams, Figure
from matplotlib.axes import Axes, subplot_class_factory
from matplotlib.patches import PathPatch, Rectangle, Arc
from matplotlib.path import Path
from matplotlib.widgets import RectangleSelector
from matplotlib.transforms import Bbox, offset_copy, IdentityTransform, \
     Affine2D
from matplotlib import cm as mpl_colormap
from matplotlib import colors as mpl_colors
from matplotlib.colorbar import Colorbar
from matplotlib.collections import RegularPolyCollection, LineCollection, \
     PatchCollection
from matplotlib.lines import Line2D
try:
    from matplotlib.offsetbox import OffsetImage, AnnotationBbox
except ImportError:
    pass
from matplotlib._png import read_png
from matplotlib.ticker import MaxNLocator, FuncFormatter, NullLocator
from mpl_toolkits.axes_grid.anchored_artists import AnchoredText
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import symbols, colors
import hardcopy as HC
try:
    import Image
except ImportError:
    from PIL import Image

#matplotlib.rcParams['path.simplify'] = False

_tango = colors.tango()
class TreeFigure(object):
    """
    Window for showing a single tree, optionally with split overview
    and detail panes.

    The navigation toolbar at the bottom is provided by matplotlib
    (http://matplotlib.sf.net/users/navigation_toolbar.html).  Its
    pan/zoom button and zoom-rectangle button provide different modes
    of mouse interaction with the figure.  When neither of these
    buttons are checked, the default mouse bindings are as follows:

    * button 1 drag: select nodes - retrieve by calling fig.selected
    * button 3 drag: pan view
    * scroll up/down: zoom in/out
    * scroll up/down with Control key: zoom y-axis
    * scroll up/down with Shift key: zoom x-axis
    * scroll up/down with 'd' key: pan view up/down
    * scroll up/down with 'e' key: pan view left/right
    * click on overview will center the detail pane on that region

    Default keybindings:

    * t: zoom out to full extent
    * +/-: zoom in/out

    Useful attributes and methods (assume an instance named *fig*):

    * fig.root - the root node (see [Node methods])
    * fig.highlight(s) - highlight and trace nodes with substring *s*
    * fig.zoom_clade(anc) - zoom to view node *anc* and all its descendants
    * fig.toggle_overview() - toggle visibility of the overview pane
    * fig.toggle_branchlabels() - ditto for branch labels
    * fig.toggle_leaflabels() - ditto for leaf labels
    * fig.decorate(func) - decorate the tree with a function (see
      :ref:`decorating TreeFigures <decorating-trees>`)
    """
    def __init__(self, data, name=None, scaled=True, div=0.25,
                 branchlabels=True, leaflabels=True, mark_named=True,
                 highlight_support=True, xoff=0, yoff=0,
                 overview=True, radial=False):
        self.overview = None
        self.overview_width = div
        self.dataplot = None
        self.dataplot_width = 0.25
        self.name = name
        self.scaled = scaled
        self.branchlabels = branchlabels
        self.leaflabels = leaflabels
        self.mark_named = mark_named
        self.xoff = xoff
        self.yoff = yoff
        self.radial = radial
        if radial:
            self.leaflabels = False
        self.highlighted = set()
        self.highlight_support = highlight_support
        if isinstance(data, tree.Node):
            root = data
        else:
            root = tree.read(data)
        self.root = root
        if not self.root:
            raise IOError, "cannot coerce data into tree.Node"
        self.name = self.name or root.treename
        pars = SubplotParams(
            left=0, right=1, bottom=0.05, top=1, wspace=0.01
            )
        fig = pyplot.figure(subplotpars=pars, facecolor="white")
        connect_events(fig.canvas)
        self.figure = fig
        self.initialize_subplots(overview)
        self.home()

    def initialize_subplots(self, overview=True):
        if not self.radial:
            tp = TreePlot(self.figure, 1, 2, 2, app=self, name=self.name,
                          scaled=self.scaled, branchlabels=self.branchlabels,
                          highlight_support=self.highlight_support,
                          leaflabels=self.leaflabels,
                          mark_named=self.mark_named)
            detail = self.figure.add_subplot(tp)
            detail.set_root(self.root)
            detail.plot_tree()
            self.detail = detail
            tp = OverviewTreePlot(
                self.figure, 121, app=self, scaled=self.scaled,
                branchlabels=False, leaflabels=False,
                mark_named=self.mark_named,
                highlight_support=self.highlight_support,
                target=self.detail
                )
            ov = self.figure.add_subplot(tp)
            ov.set_root(self.root)
            ov.plot_tree()
            self.overview = ov
            if not overview:
                self.toggle_overview(False)
            self.set_positions()

            if self.detail.nleaves < 50:
                self.toggle_overview(False)
        else:
            tp = RadialTreePlot(
                self.figure, 111, app=self, name=self.name,
                scaled=self.scaled, branchlabels=self.branchlabels,
                highlight_support=self.highlight_support,
                leaflabels=self.leaflabels, mark_named=self.mark_named
                )
            ax2 = self.figure.add_subplot(tp)
            ax2.set_root(self.root)
            ax2.plot_tree()
            self.detail = ax2

    def __get_selected_nodes(self):
        return list(self.detail.selected_nodes)

    def __set_selected_nodes(self, nodes):
        self.detail.select_nodes(nodes)

    def __del_selected_nodes(self):
        self.detail.select_nodes(None)

    selected = property(__get_selected_nodes,
                        __set_selected_nodes,
                        __del_selected_nodes)

    ## def selected_nodes(self):
    ##     return self.detail.selected_nodes

    @property
    def axes(self):
        return self.detail

    def add(self, data, name=None, support=70,
            branchlabels=False, leaflabels=True, mark_named=True):
        """
        Add a new tree in a new window

        Args:
            data: A node object or tree file.
            name (str): Name of the plot. Defaults to None
            branchlabels (bool): Whether or not to draw branch labels.
              Defaults to False
            leaflabels (bool): Whether or not to draw leaf labels.
              Defaults to True
        """
        newfig = MultiTreeFigure()
        ## newfig.add(self.root, name=self.name, support=self.support,
        ##            branchlabels=self.branchlabels)
        newfig.add(data, name=name, support=support,
                   branchlabels=branchlabels,
                   leaflabels=leaflabels,
                   mark_named=mark_named)
        return newfig

    def toggle_leaflabels(self):
        """
        Toggle leaf labels and redraw tree
        """
        self.leaflabels = not self.leaflabels
        self.detail.leaflabels = self.leaflabels
        self.redraw()

    def toggle_branchlabels(self):
        """
        Toggle branch labels and redraw tree
        """
        self.branchlabels = not self.branchlabels
        self.detail.branchlabels = self.branchlabels
        self.redraw()

    def toggle_overview(self, val=None):
        """
        Toggle overview
        """
        if val is None:
            if self.overview.get_visible():
                self.overview.set_visible(False)
                self.overview_width = 0.001
            else:
                self.overview.set_visible(True)
                self.overview_width = 0.25
        elif val:
            self.overview.set_visible(True)
            self.overview_width = val
        else:
            self.overview.set_visible(False)
            self.overview_width = 0.001
        self.set_positions()

    def set_scaled(self, scaled):
        """
        RR: Using this method gives the error:
            redraw takes exactly 1 argument(2 given)-CZ
        Define whether or not the tree is scaled and redraw tree

        Args:
            scaled (bool): Whether or not the tree is scaled.
        """
        for p in self.overview, self.detail:
            p.redraw(p.set_scaled(scaled))
        self.set_positions()

    def on_nodes_selected(self, treeplot):
        pass

    def picked(self, e):
        try:
            if e.mouseevent.button==1:
                s = e.artist.get_text()
                clipboard.copy(s)
                print s
                sys.stdout.flush()
        except:
            pass

    def ladderize(self, rev=False):
        """
        Ladderize and redraw the tree
        """
        self.root.ladderize(rev)
        self.redraw()

    def show(self):
        """
        Plot the figure in a new window
        """
        self.figure.show()

    def set_positions(self):
        ov = self.overview
        p = self.detail
        dp = self.dataplot
        height = 1.0-p.xoffset()
        if ov:
            box = [0, p.xoffset(), self.overview_width, height]
            ov.set_position(box)
        w = 1.0
        if ov:
            w -= self.overview_width
        if dp:
            w -= self.dataplot_width
        p.set_position([self.overview_width, p.xoffset(), w, height])
        if dp:
            box = [1.0-self.dataplot_width, p.xoffset(),
                   self.dataplot_width, height]
            dp.set_position(box)
        self.figure.canvas.draw_idle()

    ## def div(self, v=0.3):
    ##     assert 0 <= v < 1
    ##     self.overview_width = v
    ##     self.set_positions()
    ##     self.figure.canvas.draw_idle()

    def add_dataplot(self):
        """
        Add new plot to the side of existing plot
        """
        np = 3 if self.overview else 2
        if self.dataplot:
            self.figure.delaxes(self.dataplot)
        self.dataplot = self.figure.add_subplot(1, np, np, sharey=self.detail)
        # left, bottom, width, height (proportions)
        dleft, dbottom, dwidth, dheight = self.detail.get_position().bounds
        # give the dataplot one-quarter the width of the detail axes
        w = dwidth * 0.25
        self.detail.set_position([dleft, dbottom, dwidth-w, dheight])
        self.dataplot.set_position([1-w, dbottom, w, dheight])
        self.dataplot.xaxis.set_visible(False)
        self.dataplot.yaxis.set_visible(False)
        for x in self.dataplot.spines.values():
            x.set_visible(False)
        self.figure.canvas.draw_idle()
        return self.dataplot

    def redraw(self):
        """
        Replot the figure and overview
        """
        self.detail.redraw()
        if self.overview: self.overview.redraw()
        self.highlight()
        self.set_positions()
        self.figure.canvas.draw_idle()

    def find(self, x):
        """
        Find nodes

        Args:
            x (str): String to search
        Returns:
            list: A list of node objects found with the Node findall() method
        """
        return self.root.findall(x)

    def hlines(self, nodes, width=5, color="red", xoff=0, yoff=0):
        """
        Highlight nodes

        Args:
            nodes (list): A list of node objects
            width (float): Width of highlighted lines. Defaults to 5
            color (str): Color of highlighted lines. Defaults to red
            xoff (float): Number of units to offset lines by. Defaults to 0
            yoff (float): Number of units to offset lines by. Defaults to 0
        """
        self.overview.hlines(nodes, width=width, color=color,
                             xoff=xoff, yoff=yoff)
        self.detail.hlines(nodes, width=width, color=color,
                           xoff=xoff, yoff=yoff)

    def highlight(self, x=None, width=5, color="red"):
        """
        Highlight nodes

        Args:
            x: Str or list of Strs or Node or list of Nodes
            width (float): Width of highlighted lines. Defaults to 5
            color (str): Color of highlighted lines. Defaults to red
        """
        if x:
            nodes = set()
            if type(x) in types.StringTypes:
                nodes = self.root.findall(x)
            elif isinstance(x, tree.Node):
                nodes = set(x)
            else:
                for n in x:
                    if type(n) in types.StringTypes:
                        found = self.root.findall(n)
                        if found:
                            nodes |= set(found)
                    elif isinstance(n, tree.Node):
                        nodes.add(n)

            self.highlighted = nodes
        else:
            self.highlighted = set()
        if self.overview:
            self.overview.highlight(self.highlighted, width=width, color=color)
        self.detail.highlight(self.highlighted, width=width, color=color)
        self.figure.canvas.draw_idle()

    def home(self):
        """
        Return plot to initial size and location.
        """
        if self.overview: self.overview.home()
        self.detail.home()

    def zoom_clade(self, x):
        """
        Zoom to fit a node *x* and all its descendants in the view.

        Args:
            x: Node or str that matches the label of a node
        """
        if not isinstance(x, tree.Node):
            x = self.root[x]
        self.detail.zoom_clade(x)

    def zoom(self, factor=0.1):
        """Zoom both axes by *factor* (relative display size)."""
        self.detail.zoom(factor, factor)
        self.figure.canvas.draw_idle()

    def zx(self, factor=0.1):
        """Zoom x axis by *factor*."""
        self.detail.zoom(factor, 0)
        self.figure.canvas.draw_idle()

    def zy(self, factor=0.1):
        """Zoom y axis by *factor*."""
        self.detail.zoom(0, factor)
        self.figure.canvas.draw_idle()

    def decorate(self, func, *args, **kwargs):
        """
        Decorate the tree.

        Args:
            func (function): A function that takes a TreePlot instance as the
              first parameter, and *args* and *kwargs* as additional
              parameters.  It adds boxes, circles, etc to the TreePlot.

        Notes:
            If *kwargs* contains the key-value pair ('store', *name*),
            then the function is stored as *name* and re-called every time
            the TreePlot is redrawn, i.e., the decoration is persistent.
            Use ``rmdec(name)`` to remove the decorator from the treeplot.
        """
        self.detail.decorate(func, *args, **kwargs)

    def rmdec(self, name):
        "Remove the decoration 'name'."
        self.detail.rmdec(name)
        ## if name in self.detail.decorators:
        ##     del self.detail.decorators[name]

    def cbar(self, node, width=6, color='blue', mrca = True):
        pass
        # self.axes.cbar(nodes = node, width = width, color = color, mrca = mrca)

    def unclutter(self, *args):
        self.detail.unclutter()

    def trace_branches(self, nodes, width=4, color="blue"):
        """
        RR: What is the difference between this and highlight? -CZ
        """
        for p in self.overview, self.detail:
            p.trace_branches(nodes, width, color)

    def plot_continuous(self, *args, **kwargs):
        self.detail.plot_continuous(*args, **kwargs)

    def hardcopy(self, fname=None, relwidth=None, leafpad=1.5):
        if not relwidth:
            bbox = self.detail.get_tightbbox(self.figure.canvas.get_renderer())
            relwidth = bbox.width/bbox.height
        f = self.detail.hardcopy(
            relwidth=relwidth,
            leafpad=leafpad
            )
        f.axes.home()
        #f.axes.set_xlim(*self.detail.get_xlim())
        #f.axes.set_ylim(*self.detail.get_ylim())
        if fname:
            f.savefig(fname)
        return f

    def select_nodes(self, nodes=None):
        """
        Select nodes on the plot

        Args:
            nodes: A node or list of ndoes
        Notes:
            If only one node is given, all of the node's ancestors are
            also selected. If a list of nodes is given (even if it has only
            one node), only the given node(s) are selected.
        """
        self.detail.select_nodes(nodes)

    def decorate(self, func, *args, **kwargs): # RR: is this repeated from above? -CZ
        self.detail.decorate(func, *args, **kwargs)

    ## def dataplot(self):
    ##     ax = self.figure.add_subplot(133, sharey=self.detail)
    ##     ax.yaxis.set_visible(False)
    ##     self.dataplot = ax
    ##     return ax

    def attach_alignment(self, aln, overview=True):
        "leaf labels expected to be sequence ids"
        from Bio.Align import MultipleSeqAlignment
        from Bio.Seq import Seq
        from Bio.SeqRecord import SeqRecord
        from Bio.Alphabet import IUPAC
        from alignment import AlignmentFigure, AlignmentPlot
        if not isinstance(aln, MultipleSeqAlignment):
            from .. import align
            aln = align.read(aln)
        d = dict([ (x.id,x) for x in aln ])
        emptyseq = Seq('-'*aln.get_alignment_length(),
                       alphabet=IUPAC.ambiguous_dna)
        aln = MultipleSeqAlignment(
            [ d.get(x.label) or SeqRecord(emptyseq, id=x.label)
              for x in self.root.leaves() ]
            )
        self.aln = aln
        p = AlignmentPlot(self.figure, 133, aln=aln, app=self,
                          sharey=self.detail, showy=False)
        self.alnplot = Storage()
        self.alnplot.detail = self.figure.add_subplot(p)
        detail = self.alnplot.detail
        detail.plot_aln()
        if overview:
            self.alnplot.overview = inset_axes(
                detail, width="30%", height="20%", loc=1
                )
            overview = self.alnplot.overview
            overview.xaxis.set_major_locator(NullLocator())
            overview.yaxis.set_major_locator(NullLocator())
            overview.imshow(
                detail.array, interpolation='nearest', aspect='auto',
                origin='lower'
                )
            rect = UpdatingRect(
                [0,0], 0, 0, facecolor='black', edgecolor='cyan', alpha=0.5
                )
            overview.zoomrect = rect
            rect.target = detail
            detail.callbacks.connect('xlim_changed', rect)
            detail.callbacks.connect('ylim_changed', rect)
            overview.add_patch(rect)
            rect(overview)
        self.toggle_overview(False)
        xoff = self.detail.xoffset()
        self.detail.set_position([0, xoff, 0.3, 1.0-xoff])
        p.set_position([0.3, xoff, 0.7, 1.0-xoff])

class MultiTreeFigure(object):
    """
    Window for showing multiple trees side-by-side.

    TODO: document this
    """
    def __init__(self, trees=None, name=None, support=70,
                 scaled=True, branchlabels=False, radial=False):
        """
        *trees* are assumed to be objects suitable for passing to
        ivy.tree.read()
        """
        self.root = []
        self.name = name
        self.name2plot = {}
        self.plot = []
        self.scaled = scaled
        self.branchlabels = branchlabels
        self.radial = radial
        self.highlighted = set()
        self.divs = []
        pars = SubplotParams(
            left=0, right=1, bottom=0.05, top=1, wspace=0.04
            )
        fig = pyplot.figure(subplotpars=pars)
        connect_events(fig.canvas)
        self.figure = fig

        for x in trees or []:
            self.add(x, support=support, scaled=scaled,
                     branchlabels=branchlabels)

    def on_nodes_selected(self, treeplot):
        pass

    def clear(self):
        self.root = []
        self.name2plot = {}
        self.highlighted = set()
        self.divs = []
        self.figure.clf()

    def picked(self, e):
        try:
            if e.mouseevent.button==1:
                print e.artist.get_text()
                sys.stdout.flush()
        except:
            pass

    def getplot(self, x):
        p = None
        try:
            i = self.root.index(x)
            return self.plot[i]
        except ValueError:
            return self.name2plot.get(x)

    def add(self, data, name=None, support=70, scaled=True,
            branchlabels=False, leaflabels=True, mark_named=True):
        root = None
        if isinstance(data, tree.Node):
            root = data
        else:
            root = tree.read(data)
        if not root:
            raise IOError, "cannot coerce data into tree.Node"

        name = name or root.treename
        self.root.append(root)

        fig = self.figure
        N = len(self.plot)+1
        for i, p in enumerate(self.plot):
            p.change_geometry(1, N, i+1)
        plt = TreePlot(fig, 1, N, N, app=self, name=name, support=support,
                       scaled=scaled, branchlabels=branchlabels,
                       leaflabels=leaflabels, mark_named=mark_named)
        p = fig.add_subplot(plt)
        p.set_root(root)
        p.plot_tree()
        p.index = N-1
        self.plot.append(p)
        if name:
            assert name not in self.name2plot
            self.name2plot[name] = p

        ## global IP
        ## if IP:
        ##     def f(shell, s):
        ##         self.highlight(s)
        ##         return sorted([ x.label for x in self.highlighted ])
        ##     IP.expose_magic("highlight", f)
        ##     def f(shell, s):
        ##         self.root.ladderize()
        ##         self.redraw()
        ##     IP.expose_magic("ladderize", f)
        ##     def f(shell, s):
        ##         self.show()
        ##     IP.expose_magic("show", f)
        ##     def f(shell, s):
        ##         self.redraw()
        ##     IP.expose_magic("redraw", f)
        return p

    def show(self):
        self.figure.show()

    def redraw(self):
        for p in self.plot:
            p.redraw()
        self.figure.canvas.draw_idle()

    def ladderize(self, reverse=False):
        for n in self.root:
            n.ladderize(reverse)
        self.redraw()

    def highlight(self, s=None, add=False, width=5, color="red"):
        """
        Highlight nodes

        Args:
            s: Str or list of Strs or Node or list of Nodes
            add (bool): Whether to add to existing highlighted nodes or
              overwrite them.
            width (float): Width of highlighted lines. Defaults to 5
            color (str): Color of highlighted lines. Defaults to red
        """
        if not s:
            self.highlighted = set()
        if not add:
            self.highlighted = set()

        nodesets = [ p.root.findall(s) for p in self.plot ]

        for nodes, p in zip(nodesets, self.plot):
            if nodes:
                p.highlight(nodes, width=width, color=color)
            else:
                p.highlight()

        self.highlighted = nodesets
        self.figure.canvas.draw_idle()

        ##     for root in self.root:
        ##         for node in root.iternodes():
        ##             if node.label and (s in node.label):
        ##                 self.highlighted.add(node)
        ## self.highlight()

    def home(self):
        for p in self.plot: p.home()

def connect_events(canvas):
    mpl_connect = canvas.mpl_connect
    mpl_connect("button_press_event", onclick)
    mpl_connect("button_release_event", onbuttonrelease)
    mpl_connect("scroll_event", onscroll)
    mpl_connect("pick_event", onpick)
    mpl_connect("motion_notify_event", ondrag)
    mpl_connect("key_press_event", onkeypress)
    mpl_connect("axes_enter_event", axes_enter)
    mpl_connect("axes_leave_event", axes_leave)

class UpdatingRect(Rectangle):
    def __call__(self, p):
        self.set_bounds(*p.viewLim.bounds)
        p.figure.canvas.draw_idle()

class Tree(Axes):
    """
    matplotlib.axes.Axes subclass for rendering trees.
    """
    def __init__(self, fig, rect, *args, **kwargs):
        self.root = None
        self.app = kwargs.pop("app", None)
        self.support = kwargs.pop("support", 70.0)
        self.scaled = kwargs.pop("scaled", True)
        self.leaflabels = kwargs.pop("leaflabels", True)
        self.branchlabels = kwargs.pop("branchlabels", True)
        self._mark_named = kwargs.pop("mark_named", True)
        self.name = None
        self.leaf_fontsize = kwargs.pop("leaf_fontsize", 10)
        self.branch_fontsize = kwargs.pop("branch_fontsize", 10)
        self.branch_width = kwargs.pop("branch_width", 1)
        self.branch_color = kwargs.pop("branch_color", "black")
        self.interactive = kwargs.pop("interactive", True)
        self.decorators = kwargs.pop("decorators", [])
        ## if self.decorators:
        ##     print >> sys.stderr, "got %s decorators" % len(self.decorators)
        self.xoff = kwargs.pop("xoff", 0)
        self.yoff = kwargs.pop("yoff", 0)
        self.highlight_support = kwargs.pop("highlight_support", True)
        self.smooth_xpos = kwargs.pop("smooth_xpos", 0)
        Axes.__init__(self, fig, rect, *args, **kwargs)
        self.nleaves = 0
        self.highlighted = None
        self.highlightpatch = None
        self.pan_start = None
        if not self.decorators:
            self.decorators = [
                ("__selected_nodes__", (Tree.highlight_selected_nodes, [], {}))
                ]
        self.name2dec = dict([ (x[0], i) for i, x in
                               enumerate(self.decorators) ])
        self._active = False

        if self.interactive:
            self.callbacks.connect("ylim_changed", self.draw_labels)
        self.selector = RectangleSelector(self, self.rectselect,
                                          useblit=True)
        def f(e):
            if e.button != 1: return True
            else: return RectangleSelector.ignore(self.selector, e)
        self.selector.ignore = f
        self.xoffset_value = 0.05
        self.selected_nodes = set()
        self.leaf_offset = 4
        self.leaf_valign = "center"
        self.leaf_halign = "left"
        self.branch_offset = -5
        self.branch_valign = "center"
        self.branch_halign = "right"

        self.spines["top"].set_visible(False)
        self.spines["left"].set_visible(False)
        self.spines["right"].set_visible(False)
        self.xaxis.set_ticks_position("bottom")

    def p2y(self):
        "Convert a single display point to y-units"
        transform = self.transData.inverted().transform
        return transform([0,1])[1] - transform([0,0])[1]

    def p2x(self):
        "Convert a single display point to y-units"
        transform = self.transData.inverted().transform
        return transform([0,0])[1] - transform([1,0])[1]

    def decorate(self, func, *args, **kwargs):
        """
        Decorate the tree with function *func*.  If *kwargs* contains
        the key-value pair ('store', *name*), the decorator function
        is stored in self.decorators and called upon every redraw.
        """
        name = kwargs.pop("store", None)
        if name:
            if name in self.name2dec:
                i = self.name2dec[name]
                self.decorators[i] = (name, (func, args, kwargs))
            else:
                self.decorators.append((name, (func, args, kwargs)))
                self.name2dec = dict([ (x[0], i) for i, x in
                                       enumerate(self.decorators) ])

        func(self, *args, **kwargs)

    def rmdec(self, name):
        if name in self.name2dec:
            i = self.name2dec[name]
            del self.decorators[i]
            self.name2dec = dict([ (x[0], i) for i, x in
                                   enumerate(self.decorators) ])

    def flip(self):
        """
        Reverse the direction of the x-axis.
        """
        self.leaf_offset *= -1
        self.branch_offset *= -1
        ha = self.leaf_halign
        self.leaf_halign = "right" if ha == "left" else "left"
        ha = self.branch_halign
        self.branch_halign = "right" if ha == "left" else "left"
        self.invert_xaxis()
        self.redraw()

    def xoffset(self):
        """Space below x axis to show tick labels."""
        if self.scaled:
            return self.xoffset_value
        else:
            return 0

    def save_newick(self, filename):
        """
        Save tree as a newick file.

        Args:
            filename (str): Path to file.

        """
        if os.path.exists(filename):
            s = raw_input("File %s exists, enter 'y' to overwrite ").strip()
            if (s and s.lower() != 'y') or (not s):
                return
        import newick
        f = file(filename, "w")
        f.write(newick.string(self.root))
        f.close()

    def set_scaled(self, scaled):
        flag = self.scaled != scaled
        self.scaled = scaled
        return flag

    def cbar(self, nodes, color=None, label=None, x=None, width=8, xoff=10,
             showlabel=True, mrca=True):
        """
        Draw a 'clade' bar (i.e., along the y-axis) indicating a
        clade.  *nodes* are assumed to be one or more nodes in the
        tree.  If just one, it should be the internal node
        representing the clade of interest; otherwise, the clade of
        interest is the most recent common ancestor of the specified
        nodes.  *label* is an optional string to be drawn next to the
        bar, *offset* by the specified number of display units.  If
        *label* is ``None`` then the clade's label is used instead.

        Args:
            nodes: Node or list of nodes
            color (str): Color of the bar. Optional, defaults to None.
            label (str): Optional label for bar. If None, the clade's
              label is used instead. Defaults to None.
            width (float): Width of bar
            xoff (float): Offset from label to bar
            showlabel (bool): Whether or not to draw the label
            mrca: RR: Not quite sure what this does -CZ

        """
        xlim = self.get_xlim(); ylim = self.get_ylim()
        if color is None: color = _tango.next()
        transform = self.transData.inverted().transform

        if mrca:
            if isinstance(nodes, tree.Node):
                spec = nodes
            elif type(nodes) in types.StringTypes:
                spec = self.root.get(nodes)
            else:
                spec = self.root.mrca(nodes)

            assert spec in self.root
            label = label or spec.label
            leaves = spec.leaves()

        else:
            leaves = nodes

        n2c = self.n2c

        y = sorted([ n2c[n].y for n in leaves ])
        ymin = y[0]; ymax = y[-1]; y = (ymax+ymin)*0.5

        if x is None:
            x = max([ n2c[n].x for n in leaves ])
            _x = 0
            for lf in leaves:
                txt = self.node2label.get(lf)
                if txt and txt.get_visible():
                    _x = max(_x, transform(txt.get_window_extent())[1,0])
            if _x > x: x = _x

        v = sorted(list(transform(((0,0),(xoff,0)))[:,0]))
        xoff = v[1]-v[0]
        x += xoff

        Axes.plot(self, [x,x], [ymin, ymax], '-', linewidth=width, color=color)

        if showlabel and label:
            xo = self.leaf_offset
            if xo > 0:
                xo += width*0.5
            else:
                xo -= width*0.5
            txt = self.annotate(
                label,
                xy=(x, y),
                xytext=(xo, 0),
                textcoords="offset points",
                verticalalignment=self.leaf_valign,
                horizontalalignment=self.leaf_halign,
                fontsize=self.leaf_fontsize,
                clip_on=True,
                picker=False
                )

        self.set_xlim(xlim); self.set_ylim(ylim)

    def anctrace(self, anc, descendants=None, width=4, color="blue"):
        """
        RR: This function gives me a 'list index out of range' error
        when I try to use it -CZ
        """
        if not descendants:
            descendants = anc.leaves()
        else:
            for d in descendants:
                assert d in anc

        nodes = []
        for d in descendants:
            v = d.rootpath(anc)
            if v:
                nodes.extend(v)
        nodes = set(nodes)
        nodes.remove(anc)
        self.trace_branches(nodes, width, color)

    def trace_branches(self, nodes, width=4, color="blue"):
        n2c = self.n2c
        M = Path.MOVETO; L = Path.LINETO
        verts = []
        codes = []
        for c, pc in [ (n2c[x], n2c[x.parent]) for x in nodes
                       if (x in n2c) and x.parent ]:
            x = c.x; y = c.y
            px = pc.x; py = pc.y
            verts.append((x, y)); codes.append(M)
            verts.append((px, y)); codes.append(L)
            verts.append((px, py)); codes.append(L)
        px, py = verts[-1]
        verts.append((px, py)); codes.append(M)

        p = PathPatch(Path(verts, codes), fill=False,
                      linewidth=width, edgecolor=color)
        self.add_patch(p)
        self.figure.canvas.draw_idle()
        return p

    def highlight_selected_nodes(self, color="green"):
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        get = self.n2c.get
        coords = filter(None, [ get(n) for n in self.selected_nodes ])
        x = [ c.x for c in coords ]
        y = [ c.y for c in coords ]
        if x and y:
            self.__selected_highlight_patch = self.scatter(x, y, s=60, c=color,
                                                           zorder=100)
        self.set_xlim(xlim)
        self.set_ylim(ylim)
        self.figure.canvas.draw_idle()

    def select_nodes(self, nodes=None, add=False):
        try:
            self.__selected_highlight_patch.remove()
            self.figure.canvas.draw_idle()
        except:
            pass
        if add:
            if nodes:
                self.selected_nodes = self.selected_nodes | nodes
            if hasattr(self, "app") and self.app:
                self.app.on_nodes_selected(self)
            self.highlight_selected_nodes()
        else:
            if nodes:
                self.selected_nodes = nodes
                if hasattr(self, "app") and self.app:
                    self.app.on_nodes_selected(self)
                self.highlight_selected_nodes()
            else:
                self.selected_nodes = set()

    def rectselect(self, e0, e1):
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        s = set()
        x0, x1 = sorted((e0.xdata, e1.xdata))
        y0, y1 = sorted((e0.ydata, e1.ydata))
        add = e0.key == 'shift'
        for n, c in self.n2c.items():
            if (x0 < c.x < x1) and (y0 < c.y < y1):
                s.add(n)
        self.select_nodes(nodes = s, add = add)
        self.set_xlim(xlim)
        self.set_ylim(ylim)
        ## if s:
        ##     print "Selected:"
        ##     for n in s:
        ##         print " ", n

    def picked(self, e):
        if hasattr(self, "app") and self.app:
            self.app.picked(e)

    def window2data(self, expandx=1.0, expandy=1.0):
        """
        return the data coordinates ((x0, y0),(x1, y1)) of the plot
        window, expanded by relative units of window size
        """
        bb = self.get_window_extent()
        bbx = bb.expanded(expandx, expandy)
        return self.transData.inverted().transform(bbx.get_points())

    def get_visible_nodes(self, labeled_only=False):
        ## transform = self.transData.inverted().transform
        ## bb = self.get_window_extent()
        ## bbx = bb.expanded(1.1,1.1)
        ## ((x0, y0),(x1, y1)) = transform(bbx.get_points())
        ((x0, y0),(x1, y1)) = self.window2data(1.1, 1.1)
        #print "visible_nodes points", x0, x1, y0, y1

        if labeled_only:
            def f(v): return (y0 < v[0] < y1) and (v[2] in self.node2label)
        else:
            def f(v): return (y0 < v[0] < y1)
        for y, x, n in filter(f, self.coords):
            yield (n, x, y)

    def zoom_cxy(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view, with a fixed data point (cx, cy)
        """
        transform = self.transData.inverted().transform
        xlim = self.get_xlim(); xmid = sum(xlim)*0.5
        ylim = self.get_ylim(); ymid = sum(ylim)*0.5
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = xmid-x0; deltay = ymid-y0
        cx = cx or xmid; cy = cy or ymid
        xoff = (cx-xmid)*x
        self.set_xlim(xmid-deltax+xoff, xmid+deltax+xoff)
        yoff = (cy-ymid)*y
        self.set_ylim(ymid-deltay+yoff, ymid+deltay+yoff)
        self.adjust_xspine()

    def zoom(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view.
        """
        # get the function to convert display coordinates to data
        # coordinates
        transform = self.transData.inverted().transform
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = x0 - xlim[0]; deltay = y0 - ylim[0]
        self.set_xlim(xlim[0]+deltax, xlim[1]-deltax)
        self.set_ylim(ylim[0]+deltay, ylim[1]-deltay)
        self.adjust_xspine()

    def center_y(self, y):
        """
        Center the y-axis of the canvas on the given y value
        """
        ymin, ymax = self.get_ylim()
        yoff = (ymax - ymin) * 0.5
        self.set_ylim(y-yoff, y+yoff)
        self.adjust_xspine()

    def center_x(self, x, offset=0.3):
        """
        Center the x-axis of the canvas on the given x value
        """
        xmin, xmax = self.get_xlim()
        xspan = xmax - xmin
        xoff = xspan*0.5 + xspan*offset
        self.set_xlim(x-xoff, x+xoff)
        self.adjust_xspine()

    def center_node(self, node):
        """
        Center the canvas on the given node
        """
        c = self.n2c[node]
        y = c.y
        self.center_y(y)
        x = c.x
        self.center_x(x, 0.2)

    def do_highlight_support(self):
        """
        TODO: reconfigure this, insert into self.decorators
        """
        if self.support:
            lim = float(self.support)

        M = Path.MOVETO; L = Path.LINETO

        verts = []; codes = []
        segments = []
        def f(n):
            if n.isleaf or not n.parent: return False
            try: return float(n.label) >= lim
            except:
                try: return float(n.support) >= lim
                except: pass
            return False

        for node, coords in [ x for x in self.n2c.items() if f(x[0]) ]:
            x = coords.x; y = coords.y
            p = node.parent
            pcoords = self.n2c[p]
            px = pcoords.x; py = y
            if self.app and self.app.radial:
                pc = self.n2c[node.parent]; theta2 = pc.angle
                px = math.cos(math.radians(coords.angle))*pc.depth
                py = math.sin(math.radians(coords.angle))*pc.depth

            ## segments.append([(x, y),(px, y)])
            verts.append((x,y)); codes.append(M)
            verts.append((px,py)); codes.append(L)

        if verts:
            patch = PathPatch(Path(verts, codes), fill=False,
                              linewidth=3, edgecolor='black')
            self.add_patch(patch)

            ## self.add_artist(Line2D(
            ##     [x,px], [y,py], lw=3, solid_capstyle="butt", color="black"
            ##     ))

    def hl(self, s):
        nodes = self.root.findall(s)
        if nodes:
            self.highlight(nodes)

    def hlines(self, nodes, width=5, color="red", xoff=0, yoff=0):
        offset = IdentityTransform()
        segs = []; w = []; o = []
        for n in filter(lambda x:x.parent, nodes):
            c = self.n2c[n]; p = self.n2c[n.parent]
            segs.append(((p.x,c.y),(c.x,c.y)))
            w.append(width); o.append((xoff,yoff))
        lc = LineCollection(segs, linewidths=w, transOffset=offset, offsets=o)
        lc.set_color(color)
        Axes.add_collection(self, lc)
        ## self.drawstack.append(("hlines", [nodes], dict(width=width,
        ##                                                color=color,
        ##                                                xoff=xoff,
        ##                                                yoff=yoff)))
        self.figure.canvas.draw_idle()
        return lc

    def hardcopy(self, relwidth=0.5, leafpad=1.5):
        p = HC.TreeFigure(self.root, relwidth=relwidth, leafpad=leafpad,
                          name=self.name, support=self.support,
                          leaf_fontsize=self.leaf_fontsize,
                          branch_fontsize=self.branch_fontsize,
                          branch_width=self.branch_width,
                          branch_color=self.branch_color,
                          highlight_support=self.highlight_support,
                          branchlabels=self.branchlabels,
                          decorators=self.decorators,
                          leaflabels=self.leaflabels,
                          mark_named=self._mark_named,
                          xlim=self.get_xlim(),
                          ylim=self.get_ylim())
        return p

    def highlight(self, nodes=None, width=5, color="red"):
        if self.highlightpatch:
            try:
                self.highlightpatch.remove()
            except:
                pass
        if not nodes:
            return

        if len(nodes)>1:
            mrca = self.root.mrca(nodes)
            if not mrca:
                return
        else:
            mrca = list(nodes)[0]

        M = Path.MOVETO; L = Path.LINETO
        verts = []
        codes = []
        seen = set()
        for node, coords in [ x for x in self.n2c.items() if x[0] in nodes ]:
            x = coords.x; y = coords.y
            p = node.parent
            while p:
                pcoords = self.n2c[p]
                px = pcoords.x; py = pcoords.y
                if node not in seen:
                    verts.append((x, y)); codes.append(M)
                    verts.append((px, y)); codes.append(L)
                    verts.append((px, py)); codes.append(L)
                    seen.add(node)
                if p == mrca or node == mrca:
                    break
                node = p
                coords = self.n2c[node]
                x = coords.x; y = coords.y
                p = node.parent
        px, py = verts[-1]
        verts.append((px, py)); codes.append(M)

        self.highlightpath = Path(verts, codes)
        self.highlightpatch = PathPatch(
            self.highlightpath, fill=False, linewidth=width, edgecolor=color,
            capstyle='round', joinstyle='round'
            )
        return self.add_patch(self.highlightpatch)

    def find(self, s):
        """
        Find node(s) matching pattern s and zoom to node(s)
        """
        nodes = list(self.root.find(s))
        if nodes:
            self.zoom_nodes(nodes)

    def zoom_nodes(self, nodes, border=1.2):
        y0, y1 = self.get_ylim(); x0, x1 = self.get_xlim()
        y0 = max(0, y0); y1 = min(1, y1)

        n2c = self.n2c
        v = [ n2c[n] for n in nodes ]
        ymin = min([ c.y for c in v ])
        ymax = max([ c.y for c in v ])
        xmin = min([ c.x for c in v ])
        xmax = max([ c.x for c in v ])
        bb = Bbox(((xmin,ymin), (xmax, ymax)))

        # convert data coordinates to display coordinates
        transform = self.transData.transform
        disp_bb = [Bbox(transform(bb))]
        for n in nodes:
            if n.isleaf:
                txt = self.node2label[n]
                if txt.get_visible():
                    disp_bb.append(txt.get_window_extent())

        disp_bb = Bbox.union(disp_bb).expanded(border, border)

        # convert back to data coordinates
        points = self.transData.inverted().transform(disp_bb)
        x0, x1 = points[:,0]
        y0, y1 = points[:,1]
        self.set_xlim(x0, x1)
        self.set_ylim(y0, y1)

    def zoom_clade(self, anc, border=1.2):
        if anc.isleaf:
            self.center_node(anc)

        else:
            self.zoom_nodes(list(anc), border)

    def draw_leaf_labels(self, *args):
        leaves = list(filter(lambda x:x[0].isleaf,
                             self.get_visible_nodes(labeled_only=True)))
        psep = self.leaf_pixelsep()
        fontsize = min(self.leaf_fontsize, max(psep, 8))
        n2l = self.node2label
        transform = self.transData.transform
        sub = operator.sub

        for n in leaves:
            n2l[n[0]].set_visible(False)

        # draw leaves
        leaves_drawn = []
        for n, x, y in leaves:
            txt = self.node2label[n]
            if not leaves_drawn:
                txt.set_visible(True)
                leaves_drawn.append(txt)
                self.figure.canvas.draw_idle()
                continue

            txt2 = leaves_drawn[-1]
            y0 = y; y1 = txt2.xy[1]
            sep = sub(*transform(([0,y0],[0,y1]))[:,1])
            if sep > fontsize:
                txt.set_visible(True)
                txt.set_size(fontsize)
                leaves_drawn.append(txt)
        self.figure.canvas.draw_idle()

        if leaves_drawn:
            leaves_drawn[0].set_size(fontsize)

        return fontsize

    def draw_labels(self, *args):
        fs = max(10, self.draw_leaf_labels())
        nodes = self.get_visible_nodes(labeled_only=True)
        ## print [ x[0].id for x in nodes ]
        branches = list(filter(lambda x:(not x[0].isleaf), nodes))
        n2l = self.node2label
        for n, x, y in branches:
            t = n2l[n]
            t.set_visible(True)
            t.set_size(fs)

    def unclutter(self, *args):
        nodes = self.get_visible_nodes(labeled_only=True)
        branches = list(filter(lambda x:(not x[0].isleaf), nodes))
        psep = self.leaf_pixelsep()
        n2l = self.node2label
        fontsize = min(self.leaf_fontsize*1.2, max(psep, self.leaf_fontsize))

        drawn = []
        for n, x, y in branches:
            txt = n2l[n]
            try:
                bb = txt.get_window_extent().expanded(2, 2)
                vis = True
                for n2 in reversed(drawn):
                    txt2 = n2l[n2]
                    if bb.overlaps(txt2.get_window_extent()):
                        txt.set_visible(False)
                        vis = False
                        self.figure.canvas.draw_idle()
                        break
                if vis:
                    txt.set_visible(True)
                    txt.set_size(fontsize)
                    self.figure.canvas.draw_idle()
                    drawn.append(n)
            except RuntimeError:
                pass
                ## txt.set_visible(True)
                ## txt.set_size(fontsize)
                ## drawn.append(n)
                ## self.figure.canvas.draw_idle()

    def leaf_pixelsep(self):
        y0, y1 = self.get_ylim()
        y0 = max(0, y0)
        y1 = min(self.nleaves, y1)
        display_points = self.transData.transform(((0, y0), (0, y1)))
        # height in pixels (visible y data extent)
        height = operator.sub(*reversed(display_points[:,1]))
        pixelsep = height/((y1-y0)/self.leaf_hsep)
        return pixelsep

    def ypp(self):
        y0, y1 = self.get_ylim()
        p0, p1 = self.transData.transform(((0, y0), (0, y1)))[:,1]
        return (y1-y0)/float(p1-p0)

    def draw_labels_old(self, *args):
        if self.nleaves:
            y0, y1 = self.get_ylim()
            y0 = max(0, y0); y1 = min(1, y1)

            display_points = self.transData.transform(((0, y0), (0, y1)))
            # height in pixels (visible y data extent)
            height = operator.sub(*reversed(display_points[:,1]))
            pixelsep = height/((y1-y0)/self.leaf_hsep)
            fontsize = min(max(pixelsep-2, 8), 12)

            if pixelsep >= 8:
                for node, txt in self.node2label.items():
                    if node.isleaf:
                        if self.leaflabels:
                            c = self.n2c[node]
                            x = c.x; y = c.y
                            if (y0 < y < y1):
                                txt.set_size(fontsize)
                                txt.set_visible(True)
                    else:
                        if self.branchlabels:
                            c = self.n2c[node]
                            x = c.x; y = c.y
                            if (y0 < y < y1):
                                txt.set_size(fontsize)
                                txt.set_visible(True)
            elif pixelsep >= 4:
                for node, txt in self.node2label.items():
                    if node.isleaf:
                        txt.set_visible(False)
                    else:
                        if self.branchlabels:
                            c = self.n2c[node]
                            x = c.x; y = c.y
                            if (y0 < y < y1):
                                txt.set_size(fontsize)
                                txt.set_visible(True)
            else:
                for node, txt in self.node2label.items():
                    txt.set_visible(False)
            self.figure.canvas.draw_idle()

    def redraw(self, home=False, layout=True):
        """
        Replot the tree
        """
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        self.cla()
        if layout:
            self.layout()
        self.plot_tree()
        if self.interactive:
            self.callbacks.connect("ylim_changed", self.draw_labels)

        if home:
            self.home()
        else:
            self.set_xlim(*xlim)
            self.set_ylim(*ylim)

    def set_name(self, name):
        self.name = name
        if name:
            at = AnchoredText(
                self.name, loc=2, frameon=True,
                prop=dict(size=12, weight="bold")
                )
            at.patch.set_linewidth(0)
            at.patch.set_facecolor("white")
            at.patch.set_alpha(0.6)
            self.add_artist(at)
            return at

    def _path_to_parent(self, node):
        """
        For use in drawing branches
        """
        c = self.n2c[node]; x = c.x; y = c.y
        pc = self.n2c[node.parent]; px = pc.x; py = pc.y
        M = Path.MOVETO; L = Path.LINETO
        verts = [(x, y), (px, y), (px, py)]
        codes = [M, L, L]
        return verts, codes
        ## return [PathPatch(Path(verts, codes), fill=False,
        ##                   linewidth=width or self.branch_width,
        ##                   edgecolor=color or self.branch_color)]

    def layout(self):
        self.n2c = cartesian(self.root, scaled=self.scaled, yunit=1.0,
                             smooth=self.smooth_xpos)
        for c in self.n2c.values():
            c.x += self.xoff; c.y += self.yoff
        sv = sorted([
            [c.y, c.x, n] for n, c in self.n2c.items()
            ])
        self.coords = sv#numpy.array(sv)
        ## n2c = self.n2c
        ## self.node2linesegs = {}
        ## for node, coords in n2c.items():
        ##     x = coords.x; y = coords.y
        ##     v = [(x,y)]
        ##     if node.parent:
        ##         pcoords = n2c[node.parent]
        ##         px = pcoords.x; py = pcoords.y
        ##         v.append((px,y))
        ##         v.append((px,py))
        ##         self.node2linesegs[node] = v

    def set_root(self, root):
        self.root = root
        self.leaves = root.leaves()
        self.nleaves = len(self.leaves)
        self.leaf_hsep = 1.0#/float(self.nleaves)

        for n in root.descendants():
            if n.length is None:
                self.scaled=False; break
        self.layout()

    def plot_tree(self, root=None, **kwargs):
        """
        Draw branches and labels
        """
        if root and not self.root:
            self.set_root(root)

        if self.interactive: pyplot.ioff()

        if "branchlabels" in kwargs:
            self.branchlabels = kwargs["branchlabels"]
        if "leaflabels" in kwargs:
            self.leaflabels = kwargs["leaflabels"]
        self.yaxis.set_visible(False)
        self.create_branch_artists()
        self.create_label_artists()
        if self.highlight_support:
            self.do_highlight_support()
        self.mark_named()
        ## self.home()

        for k, v in self.decorators:
            func, args, kwargs = v
            func(self, *args, **kwargs)

        self.set_name(self.name)
        self.adjust_xspine()

        if self.interactive: pyplot.ion()

        labels = [ x.label for x in self.root.leaves() ]
        def fmt(x, pos=None):
            if x<0: return ""
            try: return labels[int(round(x))]
            except: pass
            return ""
        #self.yaxis.set_major_formatter(FuncFormatter(fmt))

        return self

    def clade_dimensions(self):
        n2c = self.n2c
        d = {}
        def recurse(n, n2c, d):
            v = []
            for c in n.children:
                recurse(c, n2c, d)
                if c.isleaf:
                    x, y = n2c[c].point()
                    x0 = x1 = x; y0 = y1 = y
                else:
                    x0, x1, y0, y1 = d[c]
                v.append((x0, x1, y0, y1))
            if v:
                x0 = n2c[n].x
                x1 = max([ x[1] for x in v ])
                y0 = min([ x[2] for x in v ])
                y1 = max([ x[3] for x in v ])
                d[n] = (x0, x1, y0, y1)
        recurse(self.root, n2c, d)
        return d

    def clade_height_pixels(self):
        ypp = self.ypp()
        d = self.clade_dimensions()
        h = {}
        for n, (x0, x1, y0, y1) in d.items():
            h[n] = (y1-y0)/ypp
        return h

    def _decimate_nodes(self, n=500):
        leaves = self.leaves
        nleaves = len(leaves)
        if nleaves > n:
            indices = numpy.linspace(0, nleaves-1, n).astype(int)
            leaves = [ leaves[i] for i in indices ]
            return set(list(chain.from_iterable([ list(x.rootpath())
                                                  for x in leaves ])))
        else:
            return self.root

    def create_branch_artists(self):
        """
        Use MPL Paths to draw branches
        """
        ## patches = []
        verts = []; codes = []
        for node in self.root.descendants():
            v, c = self._path_to_parent(node)
            verts.extend(v); codes.extend(c)
        self.branchpatch = PathPatch(
            Path(verts, codes), fill=False,
            linewidth=self.branch_width,
            edgecolor=self.branch_color
            )
        self.add_patch(self.branchpatch)
        ## for node in self._decimate_nodes():
        ##     if node.parent:
        ##         for p in self._path_to_parent(node):
        ##             patches.append(p)
        ## self.branch_patches = PatchCollection(patches, match_original=True)
        ## self.add_collection(self.branch_patches)

        ## print "enter: create_branch_artists"
        ## self.node2branch = {}
        ## for node, segs in self.node2linesegs.items():
        ##     line = Line2D(
        ##         [x[0] for x in segs], [x[1] for x in segs],
        ##         lw=self.branch_width, color=self.branch_color
        ##         )
        ##     line.set_visible(False)
        ##     Axes.add_artist(self, line)
        ##     self.node2branch[node] = line

        ## d = self.node2linesegs
        ## segs = [ d[n] for n in self.root if (n in d) ]

        ## dims = self.clade_dimensions(); ypp = self.ypp()
        ## def recurse(n, dims, clades, terminals):
        ##     stop = False
        ##     h = None
        ##     v = dims.get(n)
        ##     if v: h = (v[3]-v[2])/ypp
        ##     if (h and (h < 20)) or (not h):
        ##         stop = True
        ##         terminals.append(n)
        ##     if not stop:
        ##         clades.append(n)
        ##         for c in n.children:
        ##             recurse(c, dims, clades, terminals)
        ## clades = []; terminals = []
        ## recurse(self.root, dims, clades, terminals)
        ## segs = [ d[n] for n in self.root if (n in d) and (n in clades) ]
        ## for t in terminals:
        ##     if t.isleaf:
        ##         segs.append(d[t])
        ##     else:
        ##         x0, x1, y0, y1 = dims[t]
        ##         x, y = self.n2c[t].point()
        ##         px, py = self.n2c[t.parent].point()
        ##         segs.append(((px,py), (px,y), (x,y), (x1, y0), (x1,y1), (x,y)))

        ## lc = LineCollection(segs, linewidths=self.branch_width,
        ##                     colors = self.branch_color)
        ## self.branches_linecollection = Axes.add_collection(self, lc)
        ## print "leave: create_branch_artists"

    def create_label_artists(self):
        ## print "enter: create_label_artists"
        self.node2label = {}
        n2c = self.n2c
        for node, coords in n2c.items():
            x = coords.x; y = coords.y
            if node.isleaf and node.label and self.leaflabels:
                txt = self.annotate(
                    node.label,
                    xy=(x, y),
                    xytext=(self.leaf_offset, 0),
                    textcoords="offset points",
                    verticalalignment=self.leaf_valign,
                    horizontalalignment=self.leaf_halign,
                    fontsize=self.leaf_fontsize,
                    clip_on=True,
                    picker=True
                )
                txt.node = node
                txt.set_visible(False)
                self.node2label[node] = txt

            if (not node.isleaf) and node.label and self.branchlabels:
                txt = self.annotate(
                    node.label,
                    xy=(x, y),
                    xytext=(self.branch_offset,0),
                    textcoords="offset points",
                    verticalalignment=self.branch_valign,
                    horizontalalignment=self.branch_halign,
                    fontsize=self.branch_fontsize,
                    bbox=dict(fc="lightyellow", ec="none", alpha=0.8),
                    clip_on=True,
                    picker=True
                )
                ## txt.set_visible(False)
                txt.node = node
                self.node2label[node] = txt
        ## print "leave: create_label_artists"

    def adjust_xspine(self):
        v = sorted([ c.x for c in self.n2c.values() ])
        try:
            self.spines["bottom"].set_bounds(v[0],v[-1])
        except AttributeError:
            pass
        for t,n,s in self.xaxis.iter_ticks():
            if (n > v[-1]) or (n < v[0]):
                t.set_visible(False)

    def mark_named(self):
        if self._mark_named:
            n2c = self.n2c
            cv = [ c for n, c in n2c.items() if n.label and (not n.isleaf) ]
            x = [ c.x for c in cv ]
            y = [ c.y for c in cv ]
            if x and y:
                self.scatter(x, y, s=5, color='black')

    def home(self):
        td = self.transData
        trans = td.inverted().transform
        xmax = xmin = ymax = ymin = 0
        if self.node2label:
            try:
                v = [ x.get_window_extent() for x in self.node2label.values()
                      if x.get_visible() ]
                if v:
                    xmax = trans((max([ x.xmax for x in v ]),0))[0]
                    xmin = trans((min([ x.xmin for x in v ]),0))[0]
            except RuntimeError:
                pass

        v = self.n2c.values()
        ymin = min([ c.y for c in v ])
        ymax = max([ c.y for c in v ])
        xmin = min(xmin, min([ c.x for c in v ]))
        xmax = max(xmax, max([ c.x for c in v ]))
        xspan = xmax - xmin; xpad = xspan*0.05
        yspan = ymax - ymin; ypad = yspan*0.05
        self.set_xlim(xmin-xpad, xmax+xpad*2)
        self.set_ylim(ymin-ypad, ymax+ypad)
        self.adjust_xspine()

    def scroll(self, x, y):
        x0, x1 = self.get_xlim()
        y0, y1 = self.get_ylim()
        xd = (x1-x0)*x
        yd = (y1-y0)*y
        self.set_xlim(x0+xd, x1+xd)
        self.set_ylim(y0+yd, y1+yd)
        self.adjust_xspine()

    def plot_labelcolor(self, nodemap, state2color=None):
        if state2color is None:
            c = colors.tango()
            states = sorted(set(nodemap.values()))
            state2color = dict(zip(states, c))

        for node, txt in self.node2label.items():
            s = nodemap.get(node)
            if s is not None:
                c = state2color[s]
                if c:
                    txt.set_color(c)
        self.figure.canvas.draw_idle()

    def node_image(self, node, imgfile, maxdim=100, border=0):
        xoff = self.leaf_offset
        n = self.root[node]; c = self.n2c[n]; p = (c.x, c.y)
        img = Image.open(imgfile)
        if max(img.size) > maxdim:
            img.thumbnail((maxdim, maxdim))
        imgbox = OffsetImage(img)
        xycoords = self.node2label.get(node) or "data"
        if xycoords != "data": p = (1, 0.5)
        abox = AnnotationBbox(imgbox, p,
                              xybox=(xoff, 0.0),
                              xycoords=xycoords,
                              box_alignment=(0.0,0.5),
                              pad=0.0,
                              boxcoords=("offset points"))
        self.add_artist(abox)

    def plot_discrete(self, data, cmap=None, name=None,
                      xoff=10, yoff=0, size=15, legend=1):
        root = self.root
        if cmap is None:
            import ivy
            c = colors.tango()
            states = sorted(set(data.values()))
            cmap = dict(zip(states, c))
        n2c = self.n2c
        points = []; c = []
        d = dict([ (n, data.get(n)) for n in root if data.get(n) is not None ])
        for n, v in d.items():
            coord = n2c[n]
            points.append((coord.x, coord.y)); c.append(cmap[v])

        boxes = symbols.squares(self, points, c, size, xoff=xoff, yoff=yoff)

        if legend:
            handles = []; labels = []
            for v, c in sorted(cmap.items()):
                handles.append(Rectangle((0,0),0.5,1,fc=c))
                labels.append(str(v))
            self.legend(handles, labels, loc=legend)

        self.figure.canvas.draw_idle()
        return boxes

    def plot_continuous(self, data, mid=None, name=None, cmap=None,
                        size=15, colorbar=True):
        area = (size*0.5)*(size*0.5)*numpy.pi
        values = data.values()
        vmin = min(values); vmax = max(values)
        if mid is None:
            mid = (vmin+vmax)*0.5
            delta = vmax-vmin*0.5
        else:
            delta = max(abs(vmax-mid), abs(vmin-mid))
        norm = mpl_colors.Normalize(mid-delta, mid+delta)
        ## if cmap is None: cmap = mpl_colormap.binary
        if cmap is None: cmap = mpl_colormap.hot
        n2c = self.n2c
        X = numpy.array(
            [ (n2c[n].x, n2c[n].y, v) for n, v in data.items() if n in n2c ]
            )
        circles = self.scatter(
            X[:,0], X[:,1], s=area, c=X[:,2], cmap=cmap, norm=norm,
            zorder=1000
            )
        if colorbar:
            cbar = self.figure.colorbar(circles, ax=self, shrink=0.7)
            if name:
                cbar.ax.set_xlabel(name)

        self.figure.canvas.draw_idle()

class RadialTree(Tree):
    def layout(self):
        from ..layout_polar import calc_node_positions
        start = self.start if hasattr(self, 'start') else 0
        end = self.end if hasattr(self, 'end') else None
        self.n2c = calc_node_positions(self.root, scaled=self.scaled,
                                       start=start, end=end)
        sv = sorted([
            [c.y, c.x, n] for n, c in self.n2c.items()
            ])
        self.coords = sv

    ## def _path_to_parent(self, node, width=None, color=None):
    ##     c = self.n2c[node]; theta1 = c.angle; r = c.depth
    ##     M = Path.MOVETO; L = Path.LINETO
    ##     pc = self.n2c[node.parent]; theta2 = pc.angle
    ##     px1 = math.cos(math.radians(c.angle))*pc.depth
    ##     py1 = math.sin(math.radians(c.angle))*pc.depth
    ##     verts = [(c.x,c.y),(px1,py1)]; codes = [M,L]
    ##     #verts.append((pc.x,pc.y)); codes.append(L)
    ##     path = PathPatch(Path(verts, codes), fill=False,
    ##                      linewidth=width or self.branch_width,
    ##                      edgecolor=color or self.branch_color)
    ##     diam = pc.depth*2
    ##     t1, t2 = tuple(sorted((theta1,theta2)))
    ##     arc = Arc((0,0), diam, diam, theta1=t1, theta2=t2,
    ##               edgecolor=color or self.branch_color,
    ##               linewidth=width or self.branch_width)
    ##     return [path, arc]

    def _path_to_parent(self, node):
        c = self.n2c[node]; theta1 = c.angle; r = c.depth
        M = Path.MOVETO; L = Path.LINETO
        pc = self.n2c[node.parent]; theta2 = pc.angle
        px1 = math.cos(math.radians(c.angle))*pc.depth
        py1 = math.sin(math.radians(c.angle))*pc.depth
        verts = [(c.x,c.y),(px1,py1)]; codes = [M,L]
        t1, t2 = tuple(sorted((theta1,theta2)))
        diam = pc.depth*2
        arc = Arc((0,0), diam, diam, theta1=t1, theta2=t2)
        arcpath = arc.get_path()
        av = arcpath.vertices * pc.depth
        ac = arcpath.codes
        verts.extend(av.tolist())
        codes.extend(ac.tolist())
        return verts, codes

    def highlight(self, nodes=None, width=5, color="red"):
        if self.highlightpatch:
            try:
                self.highlightpatch.remove()
            except:
                pass
        if not nodes:
            return

        if len(nodes)>1:
            mrca = self.root.mrca(nodes)
            if not mrca:
                return
        else:
            mrca = list(nodes)[0]

        M = Path.MOVETO; L = Path.LINETO
        verts = []
        codes = []
        seen = set()
        patches = []
        for node, coords in [ x for x in self.n2c.items() if x[0] in nodes ]:
            x = coords.x; y = coords.y
            p = node.parent
            while p:
                pcoords = self.n2c[p]
                px = pcoords.x; py = pcoords.y
                if node not in seen:
                    v, c = self._path_to_parent(node)
                    verts.extend(v)
                    codes.extend(c)
                    seen.add(node)
                if p == mrca or node == mrca:
                    break
                node = p
                coords = self.n2c[node]
                x = coords.x; y = coords.y
                p = node.parent
        ## px, py = verts[-1]
        ## verts.append((px, py)); codes.append(M)
        self.highlightpath = Path(verts, codes)
        self.highlightpatch = PathPatch(
            self.highlightpath, fill=False, linewidth=width, edgecolor=color
            )
        self.add_patch(self.highlightpatch)
        ## self.highlight_patches = PatchCollection(patches, match_original=True)
        ## self.add_collection(self.highlight_patches)

class OverviewTree(Tree):
    def __init__(self, *args, **kwargs):
        kwargs["leaflabels"] = False
        kwargs["branchlabels"] = False
        Tree.__init__(self, *args, **kwargs)
        self.xaxis.set_visible(False)
        self.spines["bottom"].set_visible(False)
        self.add_overview_rect()

    def set_target(self, target):
        self.target = target

    def add_overview_rect(self):
        rect = UpdatingRect([0, 0], 0, 0, facecolor='black', edgecolor='red')
        rect.set_alpha(0.2)
        rect.target = self.target
        rect.set_bounds(*self.target.viewLim.bounds)
        self.zoomrect = rect
        self.add_patch(rect)
        ## if pyplot.isinteractive():
        self.target.callbacks.connect('xlim_changed', rect)
        self.target.callbacks.connect('ylim_changed', rect)

    def redraw(self):
        Tree.redraw(self)
        self.add_overview_rect()
        self.figure.canvas.draw_idle()

def axes_enter(e):
    ax = e.inaxes
    ax._active = True

def axes_leave(e):
    ax = e.inaxes
    ax._active = False

def onselect(estart, estop):
    b = estart.button
    ## print b, estart.key

def onkeypress(e):
    ax = e.inaxes
    k = e.key
    if ax and k == 't':
        ax.home()
    if ax and k == "down":
        ax.scroll(0, -0.1)
        ax.figure.canvas.draw_idle()
    if ax and k == "up":
        ax.scroll(0, 0.1)
        ax.figure.canvas.draw_idle()
    if ax and k == "left":
        ax.scroll(-0.1, 0)
        ax.figure.canvas.draw_idle()
    if ax and k == "right":
        ax.scroll(0.1, 0)
        ax.figure.canvas.draw_idle()
    if ax and k and k in '=+':
        ax.zoom(0.1,0.1)
    if ax and k == '-':
        ax.zoom(-0.1,-0.1)

def ondrag(e):
    ax = e.inaxes
    button = e.button
    if ax and button == 2:
        if not ax.pan_start:
            ax.pan_start = (e.xdata, e.ydata)
            return
        x, y = ax.pan_start
        xdelta = x - e.xdata
        ydelta = y - e.ydata
        x0, x1 = ax.get_xlim()
        xspan = x1-x0
        y0, y1 = ax.get_ylim()
        yspan = y1 - y0
        midx = (x1+x0)*0.5
        midy = (y1+y0)*0.5
        ax.set_xlim(midx+xdelta-xspan*0.5, midx+xdelta+xspan*0.5)
        ax.set_ylim(midy+ydelta-yspan*0.5, midy+ydelta+yspan*0.5)
        ax.adjust_xspine()

def onbuttonrelease(e):
    ax = e.inaxes
    button = e.button
    if button == 2:
        ## print "pan end"
        ax.pan_start = None

def onpick(e):
    ax = e.mouseevent.inaxes
    if ax:
        ax.picked(e)

def onscroll(e):
    ax = e.inaxes
    if ax:
        b = e.button
        ## print b
        k = e.key
        if k == None and b =="up":
            ax.zoom(0.1,0.1)
        if k == None and b =="down":
            ax.zoom(-0.1,-0.1)
        if k == "shift" and b == "up":
            ax.zoom_cxy(0.1, 0, e.xdata, e.ydata)
        if k == "shift" and b == "down":
            ax.zoom_cxy(-0.1, 0, e.xdata, e.ydata)
        if k == "control" and b == "up":
            ax.zoom_cxy(0, 0.1, e.xdata, e.ydata)
        if k == "control" and b == "down":
            ax.zoom_cxy(0, -0.1, e.xdata, e.ydata)
        if k == "d" and b == "up":
            ax.scroll(0, 0.1)
        if (k == "d" and b == "down"):
            ax.scroll(0, -0.1)
        if k == "c" and b == "up":
            ax.scroll(-0.1, 0)
        if k == "c" and b == "down":
            ax.scroll(0.1, 0)
        try: ax.adjust_xspine()
        except: pass
        ax.figure.canvas.draw_idle()

def onclick(e):
    ax = e.inaxes
    if ax and e.button==1 and hasattr(ax, "zoomrect") and ax.zoomrect:
        # overview clicked; reposition zoomrect
        r = ax.zoomrect
        x = e.xdata
        y = e.ydata
        arr = ax.transData.inverted().transform(r.get_extents())
        xoff = (arr[1][0]-arr[0][0])*0.5
        yoff = (arr[1][1]-arr[0][1])*0.5
        r.target.set_xlim(x-xoff,x+xoff)
        r.target.set_ylim(y-yoff,y+yoff)
        r(r.target)
        ax.figure.canvas.draw_idle()

    if ax and e.button==2:
        ## print "pan start", (e.xdata, e.ydata)
        ax.pan_start = (e.xdata, e.ydata)

def test_decorate(treeplot):
    import evolve
    data = evolve.brownian(treeplot.root)
    values = data.values()
    vmin = min(values); vmax = max(values)
    norm = mpl_colors.Normalize(vmin, vmax)
    cmap = mpl_colormap.binary
    n2c = treeplot.n2c
    X = numpy.array(
        [ (n2c[n].x, n2c[n].y, v)
          for n, v in data.items() if n in n2c ]
        )
    circles = treeplot.scatter(
        X[:,0], X[:,1], s=200, c=X[:,2], cmap=cmap, norm=norm,
        zorder=100
        )

class Decorator(object):
    def __init__(self, treeplot):
        self.plot = treeplot

class VisToggle(object):
    def __init__(self, name, treeplot=None, value=False):
        self.name = name
        self.plot = treeplot
        self.value = value

    def __nonzero__(self):
        return self.value

    def __repr__(self):
        return "%s: %s" % (self.name, self.value)

    def redraw(self):
        if self.plot:
            self.plot.redraw()

    def toggle(self):
        self.value = not self.value
        self.redraw()

    def show(self):
        if self.value == False:
            self.value = True
            self.redraw()

    def hide(self):
        if self.value == True:
            self.value = False
            self.redraw()

TreePlot = subplot_class_factory(Tree)
RadialTreePlot = subplot_class_factory(RadialTree)
OverviewTreePlot = subplot_class_factory(OverviewTree)

if __name__ == "__main__":
    import evolve
    root, data = evolve.test_brownian()
    plot_continuous(root, data, name="Brownian", mid=0.0)

"""
ivy - a phylogenetics library and visual shell
http://www.reelab.net/ivy

Copyright 2010 Richard Ree <rree@fieldmuseum.org>

Required: ipython, matplotlib, scipy, numpy
Useful: dendropy, biopython, etc.
"""
## This program is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 3 of the
## License, or (at your option) any later version.

## This program is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.

## You should have received a copy of the GNU General Public License
## along with this program. If not, see
## <http://www.gnu.org/licenses/>.

from . import tree, layout, contrasts, ages
from . import genbank, nexus, newick, storage
# from . import bipart
# import nodearray, data
# from . import treebase
# import db
# import contrib
try:
    import ltt as _ltt
    ltt = _ltt.ltt
except ImportError:
    pass

from . import align, sequtil
# from . import chars, align, sequtil
## try: import vis
## except RuntimeError: pass

"""
Calculate node ages from branch lengths.

The function of interest is `ages2lengths`
"""
from __future__ import print_function

def ages2lengths(node, node_ages, results={}):
    """
    Convert node ages to branch lengths

    Args:
        node (Node): Node object
        node_ages (dict): Dict mapping nodes to ages
    Returns:
        dict: mapping of nodes to lengths

    """
    for d in node.descendants():
        age = node_ages[d]
        if d.parent:
            parent_age = node_ages[d.parent]
            results[d] = parent_age - age
    return results

def min_ages(node, leaf_ages, results={}):
    """
    Calculate minimum ages given fixed ages in leaf_ages

    Args:
        node (Node): A node object
        leaf_ages (dict): A dict mapping leaf nodes to ages
    Returns:
        dict: mapping of nodes to ages
    """
    v = []
    for child in node.children:
        if child.label and (child.label in leaf_ages):
            age = leaf_ages[child.label]
            v.append(age)
            results[child] = age
        else:
            min_ages(child, leaf_ages, results)
            age = results[child]
            v.append(age)
    results[node] = max(v)
    return results

def smooth(node, node_ages, results={}):
    """
    adjust ages of internal nodes by smoothing
    RR: I don't actually know what this function does -CZ
    """
    if node.parent:
        parent_age = node_ages[node.parent]
        if node.children:
            max_child_age = max([ node_ages[child] for child in node.children ])
            # make the new age the average of parent and max child
            new_node_age = (parent_age + max_child_age)/2.0
            results[node] = new_node_age
        else:
            results[node] = node_ages[node]
    else:
        results[node] = node_ages[node]
    for child in node.children:
        smooth(child, node_ages, results)
    return results

if __name__ == "__main__":
    import newick, ascii

    s = "((((a,b),(c,d),(e,f)),g),h);"
    root = newick.parse(s)

    leaf_ages = {
        "a": 3,
        "b": 2,
        "c": 4,
        "d": 1,
        "e": 3,
        "f": 0.5,
        "g": 10,
        "h": 5,
        }

    ma = min_ages(root, leaf_ages)
    d = ma
    for i in range(10):
        d = smooth(root, d)
    for node, val in ages2lengths(root, d).items():
        node.length = val
    print(ascii.render(root, scaled=1))

import os
from subprocess import Popen, PIPE
from Bio import AlignIO
from Bio.Alphabet import IUPAC
try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO
from tempfile import NamedTemporaryFile

MUSCLE = "/usr/bin/muscle"

def muscle(seqs, cmd=None):
    if not cmd: cmd = MUSCLE
    assert os.path.exists(cmd)
    p = Popen([cmd], stdin=PIPE, stdout=PIPE)
    write = p.stdin.write
    for x in seqs:
        write(">%s\n%s\n" % (x.id, x.seq))
    out = p.communicate()[0]
    aln = AlignIO.read(StringIO(out), 'fasta', alphabet=IUPAC.ambiguous_dna)
    return aln

def musclep(seqs1, seqs2, cmd="/usr/bin/muscle"):
    assert os.path.exists(cmd)
    f1 = NamedTemporaryFile(); f2 = NamedTemporaryFile()
    for s, f in ((seqs1, f1), (seqs2, f2)):
        write = f.file.write
        for x in s: write(">%s\n%s\n" % (x.id, x.seq))
    f1.file.flush(); f2.file.flush()
    cmd += " -profile -in1 %s -in2 %s" % (f1.name, f2.name)
    p = Popen(cmd.split(), stdout=PIPE)
    out = p.communicate()[0]
    aln = AlignIO.read(StringIO(out), 'fasta', alphabet=IUPAC.ambiguous_dna)
    f1.file.close(); f2.file.close()
    return aln

def read(data, format=None, name=None):
    from types import StringTypes

    def strip(s):
        fname = os.path.split(s)[-1]
        head, tail = os.path.splitext(fname)
        tail = tail.lower()
        if tail in (".fasta", ".nex", ".nexus"):
            return head
        else:
            return fname

    if (not format):
        if (type(data) in StringTypes) and os.path.isfile(data):
            s = data.lower()
            if s.endswith("fasta"):
                format="fasta"
            for tail in ".nex", ".nexus":
                if s.endswith(tail):
                    format="nexus"
                    break

    if (not format):
        format = "fasta"

    if type(data) in StringTypes:
        if os.path.isfile(data):
            name = strip(data)
            with open(data) as f:
                return AlignIO.read(f, format, alphabet=IUPAC.ambiguous_dna)
        else:
            f = StringIO(data)
            return AlignIO.read(f, format, alphabet=IUPAC.ambiguous_dna)

    elif (hasattr(data, "tell") and hasattr(data, "read")):
        treename = strip(getattr(data, "name", None))
        return AlignIO.read(data, format, alphabet=IUPAC.ambiguous_dna)

    raise IOError("unable to read alignment from '%s'" % data)

def write(data, f, format='fasta'):
    AlignIO.write(data, f, format)

def find(aln, substr):
    """
    generator that yields (seqnum, pos) tuples for every position of
    ``subseq`` in `aln`
    """
    from sequtil import finditer
    N = len(substr)
    for i, rec in enumerate(aln):
        for j in finditer(rec.seq, substr):
            yield (i,j)

def find_id(aln, regexp):
    import re
    return [ (i,s) for i, s in enumerate(aln) if re.search(regexp, s.id) ]

def gapcols(aln, c='-'):
    from numpy import array
    a = array([ list(x.seq) for x in aln ])
    for i, col in enumerate(a.T):
        s = set(col==c)
        if len(s)==1 and True in s:
            yield i

from __future__ import print_function
import math
from array import array
from .layout import depth_length_preorder_traversal

class AsciiBuffer:
    def __init__(self, width, height):
        self.width = int(width)
        self.height = int(height)
        self._b = [ array('u', ' '*self.width) for line in range(self.height) ]

    def putstr(self, r, c, s):
        assert r < self.height
        assert c+len(s) <= self.width, "%s %s %s '%s'" % (self.width, r, c, s)
        c = int(c)
        self._b[r][c:c+len(s)+1] = array('u', s)

    def __str__(self):
        return "\n".join([ ''.join(b) for b in self._b ])

def sum_to_root(node, internodes=True, length=False):
    """
    Number of branches from node to root.

    Args:
        node (Node): A Node object
        RR: Do internodes and length do anything in this function? -CZ
    Returns:
        int: The number of branches from node to root.
    """
    i = 0
    n = node
    while 1:
        if not n.parent:
            break
        else:
            n = n.parent
            i += 1
    return i

## def depth_length_preorder_traversal(node):
##     if not node.parent:
##         node.depth = 0
##         node.length_to_root = 0.0
##     else:
##         p = node.parent
##         node.depth = p.depth + 1
##         node.length_to_root = p.length_to_root + (node.length or 0.0)
##     for ch in node.children:
##         depth_length_preorder_traversal(ch)

def smooth_cpos(node, n2c):
    for ch in node.children:
        smooth_cpos(ch, n2c)

    if node.parent and not node.isleaf:
        px = n2c[node.parent].c
        cx = min([ n2c[ch].c for ch in node.children ])
        dxp = n2c[node].c - px
        cxp = cx - n2c[node].c
        n2c[node].c = px + (cx - px)*0.5

def scale_cpos(node, n2c, scalef, root_offset):
    if node.parent:
        n2c[node].c = n2c[node.parent].c + (node.length * scalef)
    else:
        n2c[node].c = root_offset

    for ch in node.children:
        scale_cpos(ch, n2c, scalef, root_offset)

def set_rpos(node, n2c):
    for child in node.children:
        set_rpos(child, n2c)
        nc = n2c[node]
        if node.children:
            children = node.children
            c0 = n2c[children[0]]
            c1 = n2c[children[-1]]
            rmin = c0.r; rmax = c1.r
            nc.r = math.ceil(rmin + (rmax-rmin)/2.0)

def render(root, unitlen=3, minwidth=50, maxwidth=None, scaled=False,
           show_internal_labels=True):
    """
    Create the ascii tree to be shown with print()
    """
    n2c = depth_length_preorder_traversal(root)
    leaves = root.leaves(); nleaves = len(leaves)
    maxdepth = max([ n2c[lf].depth for lf in leaves ])
    max_labelwidth = max([ len(lf.label) for lf in leaves ]) + 1

    root_offset = 0
    if root.label and show_internal_labels:
        root_offset = len(root.label)

    width = maxdepth*unitlen + max_labelwidth + 2 + root_offset
    height = 2*nleaves - 1

    if width < minwidth:
        unitlen = math.ceil((minwidth - max_labelwidth - 2 - root_offset)/maxdepth)
        width = maxdepth*unitlen + max_labelwidth + 2 + root_offset

    print(width, height)
    buf = AsciiBuffer(width, height)

    for i, lf in enumerate(leaves):
        c = n2c[lf]
        c.c = width - max_labelwidth - 2
        c.r = i*2

    for node in root.postiter():
        nc = n2c[node]
        if node.children:
            children = node.children
            c0 = n2c[children[0]]
            c1 = n2c[children[-1]]
            rmin = c0.r; rmax = c1.r
            nc.r = math.ceil(rmin + (rmax-rmin)/2.0)
            nc.c = min([ n2c[ch].c for ch in children ]) - unitlen

    if not scaled:
        smooth_cpos(root, n2c)
    else:
        maxlen = max([ n2c[lf].length_to_root for lf in leaves ])
        scalef = (n2c[leaves[0]].c + 1 - root_offset)/maxlen
        scale_cpos(root, n2c, scalef, root_offset)

    for node in root.postiter():
        nc = n2c[node]
        if node.parent:
            pc = n2c[node.parent]
            for r in range(min([nc.r, pc.r]),
                           max([nc.r, pc.r])):
                buf.putstr(r, pc.c, ":")

            sym = getattr(nc, "hchar", "-")
            vbar = sym*math.floor(nc.c-pc.c)
            buf.putstr(nc.r, math.ceil(pc.c), vbar)

        if node.isleaf:
            buf.putstr(nc.r, nc.c+1, " "+node.label)
        else:
            if node.label and show_internal_labels:
                buf.putstr(nc.r, nc.c-len(node.label), node.label)

        buf.putstr(nc.r, nc.c, "+")

    return str(buf)

if __name__ == "__main__":
    from . import tree
    t = tree.read(
        "(foo,((bar,(dog,cat)dc)dcb,(shoe,(fly,(cow, bowwow)cowb)cbf)X)Y)Z;"
        )

    #t = tree.read("(((foo:4.6):5.6, (bar:6.5, baz:2.3):3.0):3.0);")
    #t = tree.read("(foo:4.6, (bar:6.5, baz:2.3)X:3.0)Y:3.0;")

    i = 1
    print(render(t, scaled=0, show_internal_labels=1))
    r = t.get("cat").parent
    tree.reroot(t, r)
    tp = t.parent
    tp.remove_child(t)
    c = t.children[0]
    t.remove_child(c)
    tp.add_child(c)
    print(render(r, scaled=0, show_internal_labels=1))

"""
For drawing big trees.  Calculate which clades can be 'collapsed' and
displayed with a placeholder.

TODO: test and develop this module further
"""
from storage import Storage

def autocollapse_info(node, collapsed, visible=True, info={}):
    """
    gather information to determine if a node should be collapsed

    *collapsed* is a set containing nodes that are already collapsed
    """
    if node not in info:
        s = Storage()
        info[node] = s
    else:
        s = info[node]

    if visible and (node in collapsed):
        visible = False

    nnodes = 1 # total number of nodes, including node
    # number of visible leaves
    nvisible = int((visible and node.isleaf) or (node in collapsed))
    ntips = int(node.isleaf)
    ntips_visible = int(node.isleaf and visible)
    s.has_labeled_descendant = False
    s.depth = 1

    for child in node.children:
        autocollapse_info(child, collapsed, visible, info)
        cs = info[child]
        nnodes += cs.nnodes
        nvisible += cs.nvisible
        ntips += cs.ntips
        ntips_visible += cs.ntips_visible
        if (child.label and (not child.isleaf)) \
           or (cs.has_labeled_descendant):
            s.has_labeled_descendant = True
        if cs.depth >= s.depth:
            s.depth = cs.depth+1
    s.nnodes = nnodes
    s.nvisible = nvisible
    s.ntips = ntips
    s.ntips_visible = ntips_visible
    return info

def autocollapse(root, collapsed=None, keep_visible=None, max_visible=1000):
    """
    traverse a tree and find nodes that should be collapsed in order
    to satify *max_visible*

    *collapsed* is a set object for storing collapsed nodes

    *keep_visible* is a set object of nodes that should not be placed
    in *collapsed*
    """
    collapsed = collapsed or set()
    keep_visible = keep_visible or set()
    ntries = 0
    while True:
        if ntries > 10:
            return
        info = autocollapse_info(root, collapsed)
        nvisible = info[root].nvisible
        if nvisible <= max_visible:
            return

        v = []
        for node in root.iternodes():
            s = info[node]
            if (node.label and (not node.isleaf) and node.parent and
                (node not in keep_visible)):
                w = s.nvisible/float(s.depth)
                if s.has_labeled_descendant:
                    w *= 0.25
                v.append((w, node, s))
        v.sort(); v.reverse()
        for w, node, s in v:
            if node not in keep_visible and s.nvisible < (nvisible-1):
                print(node)
                collapsed.add(node)
                nvisible -= s.nvisible
            if nvisible <= max_visible:
                break
        ntries += 1
    return collapsed

from __future__ import print_function
import numpy as np

def default_costmatrix(numstates, dtype=np.int):
    "a square array with zeroes along the diagonal, ones elsewhere"
    return np.logical_not(np.identity(numstates)).astype(float)

def minstates(v):
    "return the indices of v that equal the minimum"
    return np.nonzero(np.equal(v, min(v)))

def downpass(node, states, stepmatrix, chardata, node2dpv=None):
    if node2dpv is None:
        node2dpv = {}

    if not node.isleaf:
        for child in node.children:
            downpass(child, states, stepmatrix, chardata, node2dpv)

        dpv = np.zeros([len(states)])
        node2dpv[node] = dpv
        for i in states:
            for child in node.children:
                child_dpv = node2dpv[child]
                mincost = min([ child_dpv[j] + stepmatrix[i,j] \
                                for j in states ])
                dpv[i] += mincost

        #print node.label, node.dpv

    else:
        #print node.label, chardata[node.label]
        node2dpv[node] = stepmatrix[:,chardata[node.label]]

    return node2dpv

def uppass(node, states, stepmatrix, node2dpv, node2upm={},
           node2ancstates=None):
    parent = node.parent
    if not node.isleaf:
        if parent is None: # root
            dpv = node2dpv[node]
            upm = None
            node.mincost = min(dpv)
            node2ancstates = {node: minstates(dpv)}

        else:
            M = np.zeros(stepmatrix.shape)
            for i in states:
                sibs = [ c for c in parent.children if c is not node ]
                for j in states:
                    c = 0
                    for sib in sibs:
                        sibdpv = node2dpv[sib]
                        c += min([ sibdpv[x] + stepmatrix[j,x]
                                   for x in states ])
                    c += stepmatrix[j,i]

                    p_upm = node2upm.get(parent)
                    if p_upm is not None:
                        c += min(p_upm[j])

                    M[i,j] += c

            node2upm[node] = M

            v = node2dpv[node][:]
            for s in states:
                v[s] += min(M[s])
            node2ancstates[node] = minstates(v)

        for child in node.children:
            uppass(child, states, stepmatrix, node2dpv, node2upm,
                   node2ancstates)

    return node2ancstates

def ancstates(tree, chardata, stepmatrix):
    states = range(len(stepmatrix))
    return uppass(tree, states, stepmatrix,
                  downpass(tree, states, stepmatrix, chardata))

def _bindeltran(node, stepmatrix, node2dpv, node2deltr=None, ancstate=None):
    if node2deltr is None:
        node2deltr = {}

    dpv = node2dpv[node]
    if ancstate is not None:
        c, s = min([ (cost+stepmatrix[ancstate,i], i) \
                     for i, cost in enumerate(dpv) ])
    else:
        c, s = min([ (cost, i) for i, cost in enumerate(dpv) ])

    node2deltr[node] = s
    for child in node.children:
        _bindeltran(child, stepmatrix, node2dpv, node2deltr, s)

    return node2deltr

def binary_deltran(tree, chardata, stepmatrix):
    states = range(len(stepmatrix))
    node2dpv = downpass(tree, states, stepmatrix, chardata)
    node2deltr = _bindeltran(tree, stepmatrix, node2dpv)
    return node2deltr

if __name__ == "__main__":
    from pprint import pprint
    from ivy import tree
    root = tree.read("(a,((b,c),(d,(e,f))));")

    nstates = 4
    states = range(nstates)
    cm = default_costmatrix(nstates)
    chardata = dict(zip("abcdef", map(int, "000233")))
    dp = downpass(root, states, cm, chardata)

    for i, node in enumerate(root):
        if not node.label:
            node.label = "N%s" % i
        else:
            node.label = "%s (%s)" % (node.label, chardata[node.label])

    print(ascii.render(root))

##     nstates = 2
##     leaves = tree.leaves()
##     for leaf in leaves:
##         leaf.anc_cost_vector = chardata[leaf.label]

    pprint(
        #ancstates(root, chardata, cm)
        #uppass(root, states, cm, downpass(tree, states, cm, chardata))
        dp
        )

#!/usr/bin/env python
"""
Functions for evolving traits and trees.
"""
from __future__ import print_function

def brownian(root, sigma=1.0, init=0.0, values={}):
    """
    Recursively evolve a trait by Brownian motion up from the node
    *root*.

    * *sigma*: standard deviation of the normal random variate after
      one unit of branch length

    * *init*: initial value

    Returns: *values* - a dictionary mapping nodes to evolved values
    """
    from scipy.stats import norm
    values[root] = init
    for child in root.children:
        time = child.length
        random_step = norm.rvs(init, scale=sigma*time)
        brownian(child, sigma, random_step, values)
    return values

def test_brownian():
    """
    Evolve a trait up an example tree of primates:.

    ((((Homo:0.21,Pongo:0.21)N1:0.28,Macaca:0.49)N2:0.13,
    Ateles:0.62)N3:0.38,Galago:1.00)root;

    Returns: (*root*, *data*) - the root node and evolved data.
    """
    import newick
    root = newick.parse(
        "((((Homo:0.21,Pongo:0.21)N1:0.28,Macaca:0.49)N2:0.13,"\
        "Ateles:0.62)N3:0.38,Galago:1.00)root;"
        )
    print(root.ascii(scaled=True))
    evolved = brownian(root)
    for node in root.iternodes():
        print(node.label, evolved[node])
    return root, evolved

if __name__ == "__main__":
    test_brownian()

"""
Categorical Markov models with k states.
"""
from __future__ import print_function
import numpy, scipy, random
import scipy.linalg
import scipy.optimize
from math import log, exp
rand = random.Random()
uniform = rand.uniform; expovariate = rand.expovariate

LARGE = 10e10 # large -lnL value used to bound parameter optimization

class Q:
    def __init__(self, k=2, layout=None):
        """
        Represents a square transition matrix with k states.

        'layout' is a square (k,k) array of integers that index free
        rate parameters (values on the diagonal are ignored).  Cells
        with value 0 will have the first rate parameter, 1 the
        second, etc.
        """
        self.k = k
        self.range = range(k)
        self.offdiag = array(numpy.eye(k)==0, dtype=numpy.int)
        if layout is None:
            layout = zeros((k,k), numpy.int)
        self.layout = layout*self.offdiag

    def fill(self, rates):
        m = numpy.take(rates, self.layout)*self.offdiag
        v = m.sum(1) * -1
        for i in self.range:
            m[i,i] = v[i]
        return m

    def default_priors(self):
        p = 1.0/self.k
        return [p]*self.k

def sample_weighted(weights):
    u = uniform(0, sum(weights))
    x = 0.0
    for i, w in enumerate(weights):
        x += w
        if u < x:
            break
    return i

def conditionals(root, data, Q):
    nstates = Q.shape[0]
    states = range(nstates)
    nodes = [ x for x in root.postiter() ]
    nnodes = len(nodes)
    v = zeros((nnodes,nstates))
    n2i = {}

    for i, n in enumerate(nodes):
        n2i[n] = i
        if n.isleaf:
            state = data[n.label]
            try:
                state = int(state)
                v[i,state] = 1.0
            except ValueError:
                if state == '?' or state == '-':
                    v[i,:] += 1/float(nstates)
        else:
            Pv = [ (expm(Q*child.length)*v[n2i[child]]).sum(1)
                   for child in n.children ]
            v[i] = numpy.multiply(*Pv)
            # fossils
            state = None
            if n.label in data:
                state = int(data[n.label])
            elif n in data:
                state = int(data[n])
            if state != None:
                for s in states:
                    if s != state: v[i,s] = 0.0

    return dict([ (n, v[i]) for n,i in n2i.items() ])

def contrasts(root, data, Q):
    cond = conditionals(root, data, Q)
    d = {}
    for n in root.postiter(lambda x:x.children):
        nc = cond[n]; nc /= sum(nc)
        diff = 0.0
        for child in n.children:
            cc = cond[child]; cc /= sum(cc)
            diff += numpy.sum(numpy.abs(nc-cc))
        d[n] = diff
    return d

def lnL(root, data, Q, priors):
    d = conditionals(root, data, Q)
    return numpy.log(sum(d[root]*priors))

def optimize(root, data, Q, priors=None):
    Qfill = Q.fill
    if priors is None: priors = Q.default_priors()
    def f(params):
        if (params<0).any(): return LARGE
        return -lnL(root, data, Qfill(params), priors)

    # initial parameter values
    p = [1.0]*len(set(Q.layout.flat))

    v = scipy.optimize.fmin_powell(
        f, p, full_output=True, disp=0, callback=None
        )
    params, neglnL = v[:2]
    if neglnL == LARGE:
        raise Exception("ConvergenceError")
    return params, neglnL

def sim(root, n2p, s0, d=None):
    if d is None:
        d = {root:s0}
    for n in root.children:
        v = n2p[n][s0]
        i = sample_weighted(v)
        d[n] = i
        sim(n, n2p, i, d)
    return d

def stmap(root, states, ancstates, Q, condition_on_success):
    """
    This and its dependent functions below need testing and
    optimization.
    """
    results = []
    for n in root.descendants():
        si = ancstates[n.parent]
        sj = ancstates[n]
        v = simulate_on_branch(states, si, sj, Q, n.length,
                               condition_on_success)
        print(n, si, sj)
        if v:
            results.append(v)
        else:
            return None
    return results

def simulate_on_branch(states, si, sj, Q, brlen, condition_on_success):
    point = 0.0
    history = [(si, point)]
    if si != sj:  # condition on one change occurring
        lambd = -(Q[si,si])
        U = uniform(0.0, 1.0)
        # see appendix of Nielsen 2001, Genetics
        t = brlen - point
        newpoint = -(1.0/lambd) * log(1.0 - U*(1.0 - exp(-lambd * t)))
        newstate = draw_new_state(states, Q, si)
        history.append((newstate, newpoint))
        si = newstate; point = newpoint
    while 1:
        lambd = -(Q[si,si])
        rv = expovariate(lambd)
        newpoint = point + rv

        if newpoint <= brlen:  # state change along branch
            newstate = draw_new_state(states, Q, si)
            history.append((newstate, newpoint))
            si = newstate; point = newpoint
        else:
            history.append((si, brlen))
            break

    if si == sj or (not condition_on_success): # success
        return history

    return None

def draw_new_state(states, Q, si):
    """
    Given a rate matrix Q, a starting state si, and an ordered
    sequence of states, eg (0, 1), draw a new state sj with
    probability -(qij/qii)
    """
    Qrow = Q[si]
    qii = Qrow[si]
    qij_probs = [ (x, -(Qrow[x]/qii)) for x in states if x != si ]
    uni = uniform(0.0, 1.0)
    val = 0.0
    for sj, prob in qij_probs:
        val += prob
        if uni < val:
            return sj

def sample_ancstates(node, states, conditionals, n2p, fixed={}):
    """
    Sample ancestral states from their conditional likelihoods
    """
    ancstates = {}
    for n in node.preiter():
        if n in fixed:
            state = fixed[n]
        else:
            cond = conditionals[n]

            if n.parent:
                P = n2p[n]
                ancst = ancstates[n.parent]
                newstate_Prow = P[ancst]
                cond *= newstate_Prow

            cond /= sum(cond)

            rv = uniform(0.0, 1.0)
            v = 0.0
            for state, c in zip(states, cond):
                v += c
                if rv < v:
                    break
        ancstates[n] = state

    return ancstates

"""
Calculate independent contrasts

TODO: include utilities for transforming data, etc.
"""
from __future__ import print_function

def PIC(node, data, results={}):
    """
    Phylogenetic independent contrasts.

    Recursively calculate independent contrasts of a bifurcating node
    given a dictionary of trait values.

    Args:
        node (Node): A node object
        data (dict): Mapping of leaf names to character values

    Returns:
        dict: Mapping of internal nodes to tuples containing ancestral
              state, its variance (error), the contrast, and the
              contrasts's variance.

    TODO: modify to accommodate polytomies.
    """
    X = []; v = []
    for child in node.children:
        if child.children:
            PIC(child, data, results)
            child_results = results[child]
            X.append(child_results[0])
            v.append(child_results[1])
        else:
            X.append(data[child.label])
            v.append(child.length)

    Xi, Xj = X  # Xi - Xj is the contrast value
    vi, vj = v

    # Xk is the reconstructed state at the node
    Xk = ((1.0/vi)*Xi + (1/vj)*Xj) / (1.0/vi + 1.0/vj)

    # vk is the variance
    vk = node.length + (vi*vj)/(vi+vj)

    results[node] = (Xk, vk, Xi-Xj, vi+vj)

    return results

if __name__ == "__main__":
    import tree
    n = tree.read(
        "((((Homo:0.21,Pongo:0.21)N1:0.28,Macaca:0.49)N2:0.13,"\
        "Ateles:0.62)N3:0.38,Galago:1.00)N4:0.0;"
        )
    char1 = {
        "Homo": 4.09434,
        "Pongo": 3.61092,
        "Macaca": 2.37024,
        "Ateles": 2.02815,
        "Galago": -1.46968
        }

    for k, v in PIC(n, char1).items():
        print(k.label or k.id, v)

import re, sys, logging
from collections import defaultdict
try:
    from itertools import izip_longest, ifilter
except:
    from itertools import zip_longest as izip_longest
from Bio import Entrez, SeqIO
from Bio.Blast import NCBIWWW, NCBIXML
from . import storage

email = ""

def batch(iterable, size):
    """
    Take an iterable and return it in chunks (sub-iterables)

    Args:
        iterable: Any iterable
        size (int): Size of chunks
    Yields:
        Chunks of size `size`
    """
    args = [ iter(iterable) ]*size
    for x in izip_longest(fillvalue=None, *args):
        yield ifilter(None, x)

def extract_gbac(s):
    """
    Extract genbank accession

    Args:
        s (str): text string of genbank file
    Returns:
        list: Accession number(s)
    """
    gbac_re = re.compile(r'[A-Z]{1,2}[0-9]{4,7}')
    return gbac_re.findall(s, re.M)
    # RR: This also returns various other strings that match the pattern (eg.
    # protein ids)

def extract_gene(seq, gene):
    """
    RR: Not sure what format seq should be in -CZ
    """
    for t in "exon", "gene":
        for x in seq.features:
            if x.type == t:
                v = x.qualifiers.get("gene")
                if v == [gene]:
                    if x.sub_features:
                        s = [ seq[sf.location.start.position:
                                  sf.location.end.position]
                              for sf in x.sub_features ]
                        return reduce(lambda x,y:x+y, s)
                    else:
                        loc = x.location
                        return seq[loc.start.position-10:loc.end.position+10]

def gi2webenv(gilist):
    h = Entrez.esearch(
        db="nucleotide", term=" OR ".join(gilist), usehistory="y",
        retmax=len(gilist)
        )
    d = Entrez.read(h)
    return d["WebEnv"], d["QueryKey"]

def gi2tax(gi):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.elink(dbfrom='taxonomy', db='nucleotide', from_uid=gi,
                     LinkName='nucleotide_taxonomy')
    r = Entrez.read(h)[0]
    h.close()
    i = r['LinkSetDb'][0]['Link'][0]['Id']
    h = Entrez.efetch(db='taxonomy', id=i, retmode='xml')
    r = Entrez.read(h)[0]
    h.close()
    return r

def ac2gi(ac):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.esearch(db="nucleotide", term=ac, retmax=1)
    d = Entrez.read(h)['IdList'][0]
    h.close()
    return d

def acsum(aclist, batchsize=100):
    """
    fetch esummary info for list of accession numbers -- useful for
    getting gi and taxids
    """
    global email
    assert email, "set email!"
    Entrez.email = email
    results = {}
    for v in batch(aclist, batchsize):
        v = list(v)
        h = Entrez.esearch(
            db="nucleotide", retmax=len(v),
            term=" OR ".join([ "%s[ACCN]" % x for x in v ]),
            usehistory="y"
            )
        d = Entrez.read(h)
        h.close()
        # gis, but not in order of aclist
        gis = d['IdList']
        d = Entrez.read(Entrez.esummary(db='nucleotide', id=','.join(gis)),
                        validate=False)
        for x in d:
            ac = x['Caption']
            if ac in aclist:
                results[ac] = x
    return results

def fetch_aclist(aclist, batchsize=1000):
    global email
    assert email, "set email!"
    Entrez.email = email
    results = {}
    n = 0
    for v in batch(aclist, batchsize):
        v = list(v)
        h = Entrez.esearch(
            db="nucleotide",
            term=" OR ".join([ "%s[ACCN]" % x for x in v ]),
            usehistory="y"
            )
        d = Entrez.read(h)
        h.close()
        h = Entrez.efetch(db="nucleotide", rettype="gb", retmax=len(v),
                          webenv=d["WebEnv"], query_key=d["QueryKey"])
        seqs = SeqIO.parse(h, "genbank")
        for s in seqs:
            try:
                ac = s.annotations["accessions"][0]
                if ac in aclist:
                    results[ac] = s
            except:
                pass
        h.close()
        n += len(v)
        logging.info('fetched %s sequences', n)
    return results

def fetch_gilist(gilist, batchsize=1000):
    global email
    assert email, "set email!"
    Entrez.email = email
    results = {}
    for v in batch(gilist, batchsize):
        v = map(str, v)
        h = Entrez.epost(db="nucleotide", id=",".join(v), usehistory="y")
        d = Entrez.read(h)
        h.close()
        h = Entrez.efetch(db="nucleotide", rettype="gb", retmax=len(v),
                          webenv=d["WebEnv"], query_key=d["QueryKey"])
        seqs = SeqIO.parse(h, "genbank")
        for s in seqs:
            try:
                gi = s.annotations["gi"]
                if gi in v:
                    s.id = organism_id(s)
                    results[gi] = s
            except:
                pass
        h.close()
    return results

def organism_id(s):
    org = (s.annotations.get('organism') or '').replace('.', '')
    return '%s_%s' % (org.replace(' ','_'), s.id.split('.')[0])

def fetchseq(gi):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.efetch(db="nucleotide", id=str(gi), rettype="gb")
    s = SeqIO.read(h, 'genbank')
    s.id = organism_id(s)
    return s

def create_fastas(data, genes):
    fastas = dict([ (g, file(g+".fasta", "w")) for g in genes ])
    for label, seqs in data.items():
        for gene, s in zip(genes, seqs):
            if s and type(s) != str:
                tag = None
                try:
                    tag = "%s_%s" % (label, s.annotations["accessions"][0])
                except:
                    tag = "%s_%s" % (label, s.name)
                if tag:
                    fastas[gene].write(">%s\n%s\n" % (tag, s.seq))
            else:
                sys.stderr.write(("error: not an accession number? "
                                  "%s (%s %s)\n" % (s, label, gene)))

    for f in fastas.values(): f.close()

def merge_fastas(fnames, name="merged"):
    outfile = file(name+".phy", "w")
    gene2len = {}
    d = defaultdict(dict)
    for fn in fnames:
        gene = fn.split(".")[0]
        for rec in SeqIO.parse(file(fn), "fasta"):
            #sp = "_".join(rec.id.split("_")[:2])
            if rec.id.startswith("Pedicularis"):
                sp = rec.id.split("_")[1]
            else:
                sp = rec.id.split("_")[0]
            sp = "_".join(rec.id.split("_")[:-1])
            seq = str(rec.seq)
            d[sp][gene] = seq
            if gene not in gene2len:
                gene2len[gene] = len(seq)

    ntax = len(d)
    nchar = sum(gene2len.values())
    outfile.write("%s %s\n" % (ntax, nchar))
    genes = list(sorted(gene2len.keys()))
    for sp, data in sorted(d.items()):
        s = "".join([ (data.get(gene) or "".join(["?"]*gene2len[gene]))
                      for gene in genes ])
        outfile.write("%s  %s\n" % (sp, s))
    outfile.close()
    parts = file(name+".partitions", "w")
    i = 1
    for g in genes:
        n = gene2len[g]
        parts.write("DNA, %s = %s-%s\n" % (g, i, i+n-1))
        i += n
    parts.close()

def blast_closest(fasta, e=10):
    f = NCBIWWW.qblast("blastn", "nr", fasta, expect=e, hitlist_size=1)
    rec = NCBIXML.read(f)
    d = rec.descriptions[0]
    result = storage.Storage()
    gi = re.findall(r'gi[|]([0-9]+)', d.title) or None
    if gi: result.gi = int(gi[0])
    ac = re.findall(r'gb[|]([^|]+)', d.title) or None
    if ac: result.ac = ac[0].split(".")[0]
    result.title = d.title.split("|")[-1].strip()
    return result

def blast(query, e=10, n=100, entrez_query=""):
    f = NCBIWWW.qblast("blastn", "nr", query, expect=e, hitlist_size=n,
                       entrez_query=entrez_query)
    recs = NCBIXML.parse(f)
    return recs
    ## v = []
    ## for d in rec.descriptions:
    ##     result = Storage()
    ##     gi = re.findall(r'gi[|]([0-9]+)', d.title) or None
    ##     if gi: result.gi = int(gi[0])
    ##     ac = re.findall(r'gb[|]([^|]+)', d.title) or None
    ##     if ac: result.ac = ac[0].split(".")[0]
    ##     result.title = d.title.split("|")[-1].strip()
    ##     v.append(result)
    ## return v

def start_codons(seq):
    i = seq.find('ATG')
    while i != -1:
        yield i
        i = seq.find('ATG', i+3)

def search_taxonomy(q):
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.esearch(db="taxonomy", term=q)
    return Entrez.read(h)['IdList']

def fetchtax(taxid):
    global email
    assert email, "set email!"
    Entrez.email = email
    n = 1
    if not isinstance(taxid, int):
        # string, possibly with multiple values?
        try:
            taxid = taxid.strip()
            n = taxid.count(',') + 1
        except AttributeError:
            # iterable of values?
            try:
                n = len(taxid)
                taxid = ','.join(map(str, taxid))
            except TypeError:
                pass
    else:
        taxid = str(taxid)
    h = Entrez.efetch(db='taxonomy', id=taxid, retmode='xml', retmax=n)
    if n == 1:
        r = Entrez.read(h)[0]
    else:
        # a list of taxonomy results in same order of taxids
        r = Entrez.read(h)
    return r

__FIRST = re.compile('[^-]')
__LAST = re.compile('[-]*$')
def trimpos(rec):
    'return the positions of the first and last ungapped base'
    s = rec.seq.tostring()
    first = __FIRST.search(s).start()
    last = __LAST.search(s).start()-1
    return (first, last)

def fetch_DNA_seqs(terms, maxn=10000, batchsize=1000):
    """
    terms: sequence of search terms, quoted appropriately, with Entrez
           specifiers, e.g. ['"Mus musculus"[organism]']
    maxn: maximum number of sequences to return
    returns list of SeqRecord objects
    """
    global email
    assert email, "set email!"
    Entrez.email = email
    h = Entrez.esearch(db="nucleotide", term=" OR ".join(terms), usehistory="y")
    d = Entrez.read(h)
    env = d['WebEnv']; key = d['QueryKey']
    N = int(d['Count'])
    if maxn: N = min(N, maxn)
    logging.info('fetching %s sequences', N)
    retstart = 0
    seqs = []
    n = 0
    while n < N:
        h = Entrez.efetch(
            db="nucleotide", rettype='gb', webenv=env, query_key=key,
            retstart=retstart, retmax=batchsize
            )
        v = list(SeqIO.parse(h, "genbank"))
        n += len(v)
        logging.info('...fetched %s', n)
        seqs.extend(v)
        retstart += batchsize
    logging.info('...done')
    return seqs

def seqrec_taxid(seqrec):
    "extract the NCBI taxon id from a sequence record"
    for ft in seqrec.features:
        if ft.type == 'source':
            break
    try:
        for x in ft.qualifiers['db_xref']:
            if x.startswith('taxon:'):
                return int(x.split(':')[1])
    except:
        pass

#!/usr/bin/env ipython
# -*- coding: utf-8 -*-
"""
Adds to the interactive IPython/pylab environment
"""
import sys, os, re
from . import tree, align
from . import vis

def readtree(data, *args, **kwargs):
    return tree.read(data, *args, **kwargs)

def readaln(data, *args, **kwargs):
    return align.read(data, *args, **kwargs)

def treefig(*args, **kwargs):
    if len(args) == 1:
        fig = vis.TreeFigure(args[0], **kwargs)
    else:
        fig = vis.MultiTreeFigure(**kwargs)
        for arg in args:
            # print arg
            fig.add(arg)
    fig.show()
    return fig

def alnfig(*args, **kwargs):
    return vis.AlignmentFigure(*args, **kwargs)

def __maketree(self, s):
    words = s.split()
    treename = "root"
    fname = None
    if words:
        treename = words.pop(0)
        if words and os.path.isfile(words[0]):
            fname = words.pop(0)

    if not fname:
        ## msg = "\n".join([
        ##     "Name of tree file",
        ##     "(Try dragging one into the terminal):\n"
        ##     ])
        msg = "Enter the name of a tree file or a newick string:\n"
        fname = raw_input(msg).strip()

    quotes = ["'", '"']
    if fname and fname[0] in quotes:
        fname = fname[1:]
    if fname and fname[-1] in quotes:
        fname = fname[:-1]
    if fname:
        try:
            ## root = ivy.tree.read(fname)
            ## IPython.ipapi.get().to_user_ns({treename:root})
            cmd = "%s = ivy.tree.read('%s')" % (treename, fname)
            get_ipython().ex(cmd)
            print("Tree parsed and assigned to variable '%s'" % treename)
        except:
            print("Unable to parse tree file '%s'" % fname)
    else:
        print("Cancelled")

def __node_completer(self, event):
    symbol = event.symbol
    s = event.line
    if symbol:
        s = s[:-len(symbol)]
    quote = ""
    if s and s[-1] in ["'", '"']:
        quote = s[-1]
        s = s[:-1]
    #base = (re.findall(r'(\w+)\[\Z', s) or [None])[-1]
    base = "".join((re.findall(r'(\w+\.\w*)?(\.)?(\w+)\[\Z', s) or [""])[-1])
    ## print "symbol:", symbol
    ## print "line:", event.line
    ## print "s:", s
    ## print "quote:", quote
    ## print "base:", base
    ## print "obj:", self._ofind(base).get("obj")

    obj = None
    if base:
        obj = self._ofind(base).get("obj")
    ## print '\n'
    ## print 'base', base
    ## print 'obj', obj
    if obj and isinstance(obj, tree.Node):
        completions = ["'"]
        if quote:
            completions = sorted([ x.label for x in obj.labeled() ])
        return completions

    raise IPython.core.error.TryNext()

try:
    ## import IPython
    IP = get_ipython() #IPython.ipapi.get()
    IP.magic('matplotlib')
    if IP:
        #IP.expose_magic("maketree", __maketree)
        # IP.define_magic("maketree", __maketree)
        ## IP.set_hook(
        ##     "complete_command", __node_completer, re_key=r'\[*'
        ##     )
        IP.set_hook(
            "complete_command", __node_completer,
            re_key='.+[[]([\']|["])*\w*$'
            )

except:
    print(sys.exc_info()[0])
    sys.stderr.write("Magic commands and completers requires IPython >= 0.11\n")

## if __name__ == "__main__":
##     if len(sys.argv) > 1:
##         for fname in sys.argv[1:]:
##             if os.path.isfile(fname):
##                 execfile(fname)

"""
layout nodes in 2d space

The function of interest is `calc_node_positions` (aka nodepos)
"""
from __future__ import print_function
import numpy

class Coordinates:
    """
    Coordinates class for storing xy coordinates
    """
    def __init__(self, x=0, y=0):
        self.x = x
        self.y = y

    def __repr__(self):
        return "Coordinates(%g, %g)" % (self.x, self.y)

    def point(self):
        return (self.x, self.y)

def smooth_xpos(node, n2coords):
    """
    RR: What does smoothing do? -CZ
    """
    if not node.isleaf:
        children = node.children
        for ch in children:
            smooth_xpos(ch, n2coords)

        if node.parent:
            px = n2coords[node.parent].x
            cx = min([ n2coords[ch].x for ch in children ])
            n2coords[node].x = (px + cx)/2.0

def depth_length_preorder_traversal(node, n2coords=None, isroot=False):
    """
    Calculate node depth (root = depth 0) and length to root

    Args:
       node (Node): A node object

    Returns:
        dict: Mapping of nodes to coordinate objects. Coordinate
             objects have attributes "depth" and "length_to_root"
    """
    if n2coords is None:
        n2coords = {}
    coords = n2coords.get(node) or Coordinates()
    coords.node = node
    if (not node.parent) or isroot:
        coords.depth = 0
        coords.length_to_root = 0.0
    else:
        try:
            p = n2coords[node.parent]
            coords.depth = p.depth + 1
            coords.length_to_root = p.length_to_root + (node.length or 0.0)
        except KeyError:
            print(node.label, node.parent.label)
        except AttributeError:
            coords.depth = 0
            coords.length_to_root = 0
    n2coords[node] = coords

    for ch in node.children:
        depth_length_preorder_traversal(ch, n2coords, False)

    return n2coords

def calc_node_positions(node, width, height,
                        lpad=0, rpad=0, tpad=0, bpad=0,
                        scaled=True, smooth=True, n2coords=None):
    """
    Calculate where nodes should be positioned in 2d space for drawing a tree

    Args:
        node (Node): A (root) node
        width (float): The width of the canvas
        height (float): The height of the canvas
        lpad, rpad, tpad, bpad (float): Padding on the edges of the canvas.
                                Optional, defaults to 0.
        scaled (bool): Whether or not the tree is scaled. Optional, defaults to
                True.
        smooth (bool): Whether or not to smooth the tree. Optional, defaults to
                True.
    Returns:
        dict: Mapping of nodes to Coordinates object
    Notes:
        Origin is at upper left
    """
    width -= (lpad + rpad)
    height -= (tpad + bpad)

    if n2coords is None:
        n2coords = {}
    depth_length_preorder_traversal(node, n2coords=n2coords)
    leaves = node.leaves()
    nleaves = len(leaves)
    maxdepth = max([ n2coords[lf].depth for lf in leaves ])
    unitwidth = width/float(maxdepth)
    unitheight = height/(nleaves-1.0)

    xoff = (unitwidth * 0.5)
    yoff = (unitheight * 0.5)

    if scaled:
        maxlen = max([ n2coords[lf].length_to_root for lf in leaves ])
        scale = width/maxlen

    for i, lf in enumerate(leaves):
        c = n2coords[lf]
        c.y = i * unitheight
        if not scaled:
            c.x = width
        else:
            c.x = c.length_to_root * scale

    for n in node.postiter():
        c = n2coords[n]
        if (not n.isleaf) and n.children:
            children = n.children
            ymax = n2coords[children[0]].y
            ymin = n2coords[children[-1]].y
            c.y = (ymax + ymin)/2.0
            if not scaled:
                c.x = min([ n2coords[ch].x for ch in children ]) - unitwidth
            else:
                c.x = c.length_to_root * scale

    if (not scaled) and smooth:
        for i in range(10):
            smooth_xpos(node, n2coords)

    for coords in n2coords.values():
        coords.x += lpad
        coords.y += tpad

    for n, coords in n2coords.items():
        if n.parent:
            p = n2coords[n.parent]
            coords.px = p.x; coords.py = p.y
        else:
            coords.px = None; coords.py = None

    return n2coords

nodepos = calc_node_positions

def cartesian(node, xscale=1.0, leafspace=None, scaled=True, n2coords=None,
              smooth=0, array=numpy.array, ones=numpy.ones, yunit=None):
    """
    RR: What is the difference between this function and calc_node_positions?
        Is it being used anywhere? -CZ
    """

    if n2coords is None:
        n2coords = {}

    depth_length_preorder_traversal(node, n2coords, True)
    leaves = node.leaves()
    nleaves = len(leaves)

    # leafspace is a vector that should sum to nleaves
    if leafspace is None:
        try: leafspace = [ float(x.leafspace) for x in leaves ]
        except: leafspace = numpy.zeros((nleaves,))
    assert len(leafspace) == nleaves
    #leafspace = array(leafspace)/(sum(leafspace)/float(nleaves))

    maxdepth = max([ n2coords[lf].depth for lf in leaves ])
    depth = maxdepth * xscale
    #if not yunit: yunit = 1.0/nleaves
    yunit = 1.0

    if scaled:
        maxlen = max([ n2coords[lf].length_to_root for lf in leaves ])
        depth = maxlen

    y = 0
    for i, lf in enumerate(leaves):
        c = n2coords[lf]
        yoff = 1 + (leafspace[i] * yunit)
        c.y = y + yoff*0.5
        y += yoff
        if not scaled:
            c.x = depth
        else:
            c.x = c.length_to_root

    for n in node.postiter():
        c = n2coords[n]
        if not n.isleaf:
            children = n.children
            v = [n2coords[children[0]].y, n2coords[children[-1]].y]
            v.sort()
            ymin, ymax = v
            c.y = (ymax + ymin)/2.0
            if not scaled:
                c.x = min([ n2coords[ch].x for ch in children ]) - 1.0
            else:
                c.x = c.length_to_root

    if not scaled:
        for i in range(smooth):
            smooth_xpos(node, n2coords)

    return n2coords

if __name__ == "__main__":
    import tree
    node = tree.read("(a:3,(b:2,(c:4,d:5):1,(e:3,(f:1,g:1):2):2):2);")
    for i, n in enumerate(node.iternodes()):
        if not n.isleaf:
            n.label = "node%s" % i
    node.label = "root"
    n2c = calc_node_positions(node, width=10, height=10, scaled=True)

    from pprint import pprint
    pprint(n2c)

"""
Compute lineages through time
"""
import numpy

# RR: Should results be set to None and then defined in the function to avoid
# problems with mutable defaults in functions? -CZ
def traverse(node, t=0, results=[]):
    """
    Recursively traverse the tree and collect information about when
    nodes split and how many lineages are added by its splitting.
    """
    if node.children:
        ## if not node.label:
        ##     node.label = str(node.id)
        results.append((t, len(node.children)-1))
        for child in node.children:
            traverse(child, t+child.length, results)
    return results

def ltt(node):
    """
    Calculate lineages through time.  The tree is assumed to be an
    ultrametric chronogram (extant leaves, with branch lengths
    proportional to time).

    Args:
        node (Node): A node object. All nodes should have branch lengths.

    Returns:
        tuple: (times, diversity) - 1D-arrays containing the results.
    """
    v = traverse(node) # v is a list of (time, diversity) values
    v.sort()
    # for plotting, it is easiest if x and y values are in separate
    # sequences, so we create a transposed array from v
    times, diversity = numpy.array(v).transpose()
    return times, diversity.cumsum()

def test():
    import newick, ascii
    n = newick.parse("(((a:1,b:2):3,(c:3,d:1):1,(e:0.5,f:3):2.5):1,g:4);")
    v = ltt(n)
    print(ascii.render(n, scaled=1))
    for t, n in v:
        print(t, n)

if __name__ == "__main__":
    test()

"""
Parse newick strings.

The function of interest is `parse`, which returns the root node of
the parsed tree.
"""
from __future__ import print_function, absolute_import, division, unicode_literals
import string, sys, re, shlex, itertools
try:
    from cStringIO import StringIO
except:
    from io import StringIO

## def read(s):
##     try:
##         s = file(s).read()
##     except:
##         try:
##             s = s.read()
##         except:
##             pass
##     return parse(s)

LABELCHARS = '-.|/?#&'
META = re.compile(r'([^,=\s]+)\s*=\s*(\{[^=}]*\}|"[^"]*"|[^,]+)?')

def add_label_chars(chars):
    global LABELCHARS
    LABELCHARS += chars

class Error(Exception):
    pass

class Tokenizer(shlex.shlex):
    """Provides tokens for parsing newick strings."""
    def __init__(self, infile):
        global LABELCHARS
        shlex.shlex.__init__(self, infile, posix=False)
        self.commenters = ''
        self.wordchars = self.wordchars+LABELCHARS
        self.quotes = "'"

    def parse_embedded_comment(self):
        ws = self.whitespace
        self.whitespace = ""
        v = []
        while 1:
            token = self.get_token()
            if token == '':
                sys.stdout.write('EOF encountered mid-comment!\n')
                break
            elif token == ']':
                break
            elif token == '[':
                self.parse_embedded_comment()
            else:
                v.append(token)
        self.whitespace = ws
        return "".join(v)
        ## print "comment:", v

def parse(data, ttable=None, treename=None):
    """
    Parse a newick string.

    Args:
        data: Any file-like object that can be coerced into shlex, or
          a string (converted to StringIO)
        ttable (dict): Mapping of node labels in the newick string
          to other values.

    Returns:
        Node: The root node.
    """
    from .tree import Node

    if isinstance(data, str):
        data = StringIO(data)

    start_pos = data.tell()
    tokens = Tokenizer(data)

    node = None; root = None
    lp=0; rp=0; rooted=1

    previous = None

    ni = 0  # node id counter (preorder) - zero-based indexing
    li = 0  # leaf index counter
    ii = 0  # internal node index counter
    pi = 0  # postorder sequence
    while 1:
        token = tokens.get_token()
        #print token,
        if token == ';' or token == tokens.eof:
            assert lp == rp, \
                   "unbalanced parentheses in tree description: (%s, %s)" \
                   % (lp, rp)
            break

        # internal node
        elif token == '(':
            lp = lp+1
            newnode = Node()
            newnode.ni = ni; ni += 1
            newnode.isleaf = False
            newnode.ii = ii; ii += 1
            newnode.treename = treename
            if node:
                if node.children: newnode.left = node.children[-1].right+1
                else: newnode.left = node.left+1
                node.add_child(newnode)
            else:
                newnode.left = 1; newnode.right = 2
            newnode.right = newnode.left+1
            node = newnode

        elif token == ')':
            rp = rp+1
            node = node.parent
            node.pi = pi; pi += 1
            if node.children:
                node.right = node.children[-1].right + 1

        elif token == ',':
            node = node.parent
            if node.children:
                node.right = node.children[-1].right + 1

        # branch length
        elif token == ':':
            token = tokens.get_token()
            if token == '[':
                node.length_comment = tokens.parse_embedded_comment()
                token = tokens.get_token()

            if not (token == ''):
                try: brlen = float(token)
                except ValueError as exc:
                    raise ValueError(
                        "invalid literal for branch length, '{}'".format(token))
            else:
                raise Error('unexpected end-of-file (expecting branch length)')

            node.length = brlen
        # comment
        elif token == '[':
            node.comment = tokens.parse_embedded_comment()
            if node.comment[0] == '&':
                # metadata
                meta = META.findall(node.comment[1:])
                if meta:
                    for k, v in meta:
                        v = eval(v.replace('{','(').replace('}',')'))
                        node.meta[k] = v

        # leaf node or internal node label
        else:
            if previous != ')': # leaf node
                if ttable:
                    try:
                        ttoken = (ttable.get(int(token)) or
                                  ttable.get(token))
                    except ValueError:
                        ttoken = ttable.get(token)
                    if ttoken:
                        token = ttoken
                newnode = Node()
                newnode.ni = ni; ni += 1
                newnode.pi = pi; pi += 1
                newnode.label = "_".join(token.split()).replace("'", "")
                newnode.isleaf = True
                newnode.li = li; li += 1
                if node.children: newnode.left = node.children[-1].right+1
                else: newnode.left = node.left+1
                newnode.right = newnode.left+1
                newnode.treename = treename
                node.add_child(newnode)
                node = newnode
            else: # label
                if ttable:
                    node.label = ttable.get(token, token)
                else:
                    node.label = token

        previous = token
    node.isroot = True
    return node

## def string(node, length_fmt=":%s", end=True, newline=True):
##     "Recursively create a newick string from node."
##     if not node.isleaf:
##         node_str = "(%s)%s" % \
##                    (",".join([ string(child, length_fmt, False, newline) \
##                                for child in node.children ]),
##                     node.label or ""
##                     )
##     else:
##         node_str = "%s" % node.label

##     if node.length is not None:
##         length_str = length_fmt % node.length
##     else:
##         length_str = ""

##     semicolon = ""
##     if end:
##         if not newline:
##             semicolon = ";"
##         else:
##             semicolon = ";\n"
##     s = "%s%s%s" % (node_str, length_str, semicolon)
##     return s

## def from_nexus(infile, bufsize=None):
##     bufsize = bufsize or 1024*5000
##     TTABLE = re.compile(r'\btranslate\s+([^;]+);', re.I | re.M)
##     TREE = re.compile(r'\btree\s+([_.\w]+)\s*=[^(]+(\([^;]+;)', re.I | re.M)
##     s = infile.read(bufsize)
##     ttable = TTABLE.findall(s) or None
##     if ttable:
##         items = [ shlex.split(line) for line in ttable[0].split(",") ]
##         ttable = dict([ (k, v.replace(" ", "_")) for k, v in items ])
##     trees = TREE.findall(s)
##     ## for i, t in enumerate(trees):
##     ##     t = list(t)
##     ##     if ttable:
##     ##         t[1] = "".join(
##     ##             [ ttable.get(x, "_".join(x.split()).replace("'", ""))
##     ##               for x in shlex.shlex(t[1]) ]
##     ##             )
##     ##     trees[i] = t
##     ## return trees
##     return ttable, trees

def parse_ampersand_comment(s):
    import pyparsing
    pyparsing.ParserElement.enablePackrat()
    from pyparsing import Word, Literal, QuotedString, CaselessKeyword, \
         OneOrMore, Group, Optional, Suppress, Regex, Dict
    word = Word(string.ascii_letters+string.digits+"%_")
    key = word.setResultsName("key") + Suppress("=")
    single_value = (Word(string.ascii_letters+string.digits+"-.") |
                    QuotedString("'") |
                    QuotedString('"'))
    range_value = Group(Suppress("{") +
                        single_value.setResultsName("min") +
                        Suppress(",") +
                        single_value.setResultsName("max") +
                        Suppress("}"))
    pair = (key + (single_value | range_value).setResultsName("value"))
    g = OneOrMore(pair)
    d = []
    for x in g.searchString(s):
        v = x.value
        if isinstance(v, str):
            try: v = float(v)
            except ValueError: pass
        else:
            try: v = map(float, v.asList())
            except ValueError: pass
        d.append((x.key, v))
    return d

# def test_parse_comment():
#     v = (("height_median=1.1368683772161603E-13,height=9.188229043880098E-14,"
#           "height_95%_HPD={5.6843418860808015E-14,1.7053025658242404E-13},"
#           "height_range={0.0,2.8421709430404007E-13}"),
#          "R", "lnP=-154.27154502342688,lnP=-24657.14341301901",
#          'states="T-lateral"')
#     for s in v:
#         print "input:", s
#         print dict(parse_ampersand_comment(s))

from __future__ import print_function
import itertools
from . import newick

class Newick(object):
    """
    convenience class for storing the results of a newick tree
    record from a nexus file, as parsed by newick.nexus_iter
    """
    def __init__(self, parse_results=None, ttable={}):
        self.name = ""
        self.comment = ""
        self.root_comment = ""
        self.newick = ""
        self.ttable = ttable
        if parse_results: self.populate(parse_results)

    def populate(self, parse_results, ttable={}):
        self.name = parse_results.tree_name
        self.comment = parse_results.tree_comment
        self.root_comment = parse_results.root_comment
        self.newick = parse_results.newick
        if ttable: self.ttable = ttable

    def parse(self, newick=newick):
        assert self.newick
        self.root = newick.parse(
            self.newick, ttable=self.ttable, treename=self.name
            )
        return self.root

def fetchaln(fname):
    """Fetch alignment"""
    from Bio.Nexus import Nexus
    n = Nexus.Nexus(fname)
    return n

def split_blocks(infile):
    try:
        from cStringIO import StringIO
    except:
        from io import StringIO
    dropwhile = itertools.dropwhile; takewhile = itertools.takewhile
    blocks = []
    not_begin = lambda s: not s.lower().startswith("begin")
    not_end = lambda s: not s.strip().lower() in ("end;", "endblock;")
    while 1:
        f = takewhile(not_end, dropwhile(not_begin, infile))
        try:
            b = f.next().split()[-1][:-1]
            blocks.append((b, StringIO("".join(list(f)))))
        except StopIteration:
            break
    return blocks

def parse_treesblock(infile):
    import string
    from pyparsing import Optional, Word, Regex, CaselessKeyword, Suppress
    from pyparsing import QuotedString
    comment = Optional(Suppress("[&") + Regex(r'[^]]+') + Suppress("]"))
    name = Word(alphanums+"_") | QuotedString("'")
    newick = Regex(r'[^;]+;')
    tree = (CaselessKeyword("tree").suppress() +
            Optional("*").suppress() +
            name.setResultsName("tree_name") +
            comment.setResultsName("tree_comment") +
            Suppress("=") +
            comment.setResultsName("root_comment") +
            newick.setResultsName("newick"))
    ## treesblock = Group(beginblock +
    ##                    Optional(ttable.setResultsName("ttable")) +
    ##                    Group(OneOrMore(tree)) +
    ##                    endblock)

    def parse_ttable(f):
        ttable = {}
        while True:
            s = f.next().strip()
            if s.lower() == ";":
                break
            if s[-1] in ",;":
                s = s[:-1]
            k, v = s.split()
            ttable[k] = v
            if s[-1] == ";":
                break
        return ttable

    ttable = {}
    while True:
        try:
            s = infile.next().strip()
        except StopIteration:
            break
        if s.lower() == "translate":
            ttable = parse_ttable(infile)
            # print("ttable: %s" % len(ttable))
        else:
            match = tree.parseString(s)
            yield Newick(match, ttable)

def iter_trees(infile):
    import pyparsing
    pyparsing.ParserElement.enablePackrat()
    from pyparsing import (
        Word, Literal, QuotedString, CaselessKeyword, CharsNotIn,
        OneOrMore, Group, Optional, Suppress, Regex, Dict, ZeroOrMore,
        alphanums, nums)
    comment = Optional(Suppress("[&") + Regex(r'[^]]+') + Suppress("]"))
    name = Word(alphanums+"_.") | QuotedString("'")
    newick = Regex(r'[^;]+;')
    tree = (CaselessKeyword("tree").suppress() +
            Optional("*").suppress() +
            name.setResultsName("tree_name") +
            comment.setResultsName("tree_comment") +
            Suppress("=") +
            comment.setResultsName("root_comment") +
            newick.setResultsName("newick"))

    def not_begin(s):
        # print('not_begin', s)
        return s.strip().lower() != "begin trees;"
    def not_end(s):
        # print('not_end', s)
        return s.strip().lower() not in ("end;", "endblock;")
    def parse_ttable(f):
        ttable = {}
        # com = Suppress('[') + ZeroOrMore(CharsNotIn(']')) + Suppress(']')
        com = Suppress('[' + ZeroOrMore(CharsNotIn(']') + ']'))
        while True:
            s = next(f).strip()
            if not s:
                continue
            s = com.transformString(s).strip()
            if s.lower() == ";":
                break
            b = False
            if s[-1] in ",;":
                if s[-1] == ';':
                    b = True
                s = s[:-1]
            # print(s)
            k, v = s.split()
            ttable[k] = v
            if b:
                break
        return ttable

    # read lines between "begin trees;" and "end;"
    f = itertools.takewhile(not_end, itertools.dropwhile(not_begin, infile))
    s = next(f).strip().lower()
    if s != "begin trees;":
        print("Expecting 'begin trees;', got %s" % s, file=sys.stderr)
        raise StopIteration
    ttable = {}
    while True:
        try:
            s = next(f).strip()
        except StopIteration:
            break
        if not s:
            continue
        if s.lower() == "translate":
            ttable = parse_ttable(f)
            # print "ttable: %s" % len(ttable)
        elif s.split()[0].lower()=='tree':
            match = tree.parseString(s)
            yield Newick(match, ttable)

try:
    from itertools import izip, imap
except ImportError:
    izip = zip
    imap = map
import numpy

def finditer(seq, substr, start=0):
    """
    Find substrings within a sequence

    Args:
        seq (str): A sequence.
        substr (str): A subsequence to search for
        start (int): Starting index. Defaults to 0
    Yields:
        int: Starting indicies of where the substr was found in seq
    """
    N = len(substr)
    i = seq.find(substr, start)
    while i >= 0:
        yield i
        i = seq.find(substr, i+N)

def gapidx(seq, gapchar='-'):
    """
    For a sequence with gaps, calculate site positions without gaps

    Args:
        seq (list): Each element of the list is one character in a sequence.
        gapchar (str): The character gaps are coded as. Defaults to '-'
    Returns:
        array: An array where the first element corresponds to range(number of
        characters that are not gaps) and the second element is the indicies
        of all characters that are not gaps.
    """
    a = numpy.array(seq)
    idx = numpy.arange(len(a))
    nongap = idx[a != gapchar]
    return numpy.array((numpy.arange(len(nongap)), nongap))

def find_stop_codons(seq, pos=0):
    """
    Find stop codons within sequence (in reading frame)

    Args:
        seq (str): A sequence
        pos (int): Starting position. Defaults to 0.
    Yields:
        tuple: The index where the stop codon starts
        and which stop codon was found.
    """
    s = seq[pos:]
    it = iter(s)
    g = imap(lambda x:"".join(x), izip(it, it, it))
    for i, x in enumerate(g):
        if x in ("TAG", "TAA", "TGA"):
            yield pos+(i*3), x

from __future__ import print_function
from operator import itemgetter
from heapq import nlargest
from itertools import repeat
try:
    from itertools import ifilter
except:
    ifilter = filter

class Storage(dict):
    """
    A Storage object is like a dictionary except `obj.foo` can be used
    in addition to `obj['foo']`.

    From web2py/gluon/storage.py by Massimo Di Pierro (www.web2py.com)
    """
    def __getattr__(self, key):
        try:
            return self[key]
        except KeyError:
            return None

    def __setattr__(self, key, value):
        self[key] = value

    def __delattr__(self, key):
        del self[key]

    def __repr__(self):
        return '<Storage ' + dict.__repr__(self) + '>'

    def __getstate__(self):
        return dict(self)

    def __setstate__(self, value):
        for (k, v) in value.items():
            self[k] = v

class MaxDict(dict):
    def __setitem__(self, key, value):
        v = self.get(key)
        if value > v:
            dict.__setitem__(self, key, value)

#from http://code.activestate.com/recipes/576611/
class Counter(dict):
    """Dict subclass for counting hashable objects.  Sometimes called a bag
    or multiset.  Elements are stored as dictionary keys and their counts
    are stored as dictionary values.

    >>> Counter('zyzygy')
    Counter({'y': 3, 'z': 2, 'g': 1})

    """

    def __init__(self, iterable=None, **kwds):
        """Create a new, empty Counter object.  And if given, count elements
        from an input iterable.  Or, initialize the count from another mapping
        of elements to their counts.

        >>> c = Counter()                           # a new, empty counter
        >>> c = Counter('gallahad')                 # a new counter from an iterable
        >>> c = Counter({'a': 4, 'b': 2})           # a new counter from a mapping
        >>> c = Counter(a=4, b=2)                   # a new counter from keyword args

        """
        self.update(iterable, **kwds)

    def __missing__(self, key):
        return 0

    def most_common(self, n=None):
        """List the n most common elements and their counts from the most
        common to the least.  If n is None, then list all element counts.

        >>> Counter('abracadabra').most_common(3)
        [('a', 5), ('r', 2), ('b', 2)]

        """
        if n is None:
            return sorted(self.iteritems(), key=itemgetter(1), reverse=True)
        return nlargest(n, self.iteritems(), key=itemgetter(1))

    def elements(self):
        """Iterator over elements repeating each as many times as its count.

        >>> c = Counter('ABCABC')
        >>> sorted(c.elements())
        ['A', 'A', 'B', 'B', 'C', 'C']

        If an element's count has been set to zero or is a negative number,
        elements() will ignore it.

        """
        for elem, count in self.iteritems():
            for _ in repeat(None, count):
                yield elem

    # Override dict methods where the meaning changes for Counter objects.

    @classmethod
    def fromkeys(cls, iterable, v=None):
        raise NotImplementedError(
            'Counter.fromkeys() is undefined.  Use Counter(iterable) instead.')

    def update(self, iterable=None, **kwds):
        """Like dict.update() but add counts instead of replacing them.

        Source can be an iterable, a dictionary, or another Counter instance.

        >>> c = Counter('which')
        >>> c.update('witch')           # add elements from another iterable
        >>> d = Counter('watch')
        >>> c.update(d)                 # add elements from another counter
        >>> c['h']                      # four 'h' in which, witch, and watch
        4

        """
        if iterable is not None:
            if hasattr(iterable, 'iteritems'):
                if self:
                    self_get = self.get
                    for elem, count in iterable.iteritems():
                        self[elem] = self_get(elem, 0) + count
                else:
                    dict.update(self, iterable) # fast path when counter is empty
            else:
                self_get = self.get
                for elem in iterable:
                    self[elem] = self_get(elem, 0) + 1
        if kwds:
            self.update(kwds)

    def copy(self):
        'Like dict.copy() but returns a Counter instance instead of a dict.'
        return Counter(self)

    def __delitem__(self, elem):
        'Like dict.__delitem__() but does not raise KeyError for missing values.'
        if elem in self:
            dict.__delitem__(self, elem)

    def __repr__(self):
        if not self:
            return '%s()' % self.__class__.__name__
        items = ', '.join(map('%r: %r'.__mod__, self.most_common()))
        return '%s({%s})' % (self.__class__.__name__, items)

    # Multiset-style mathematical operations discussed in:
    #       Knuth TAOCP Volume II section 4.6.3 exercise 19
    #       and at http://en.wikipedia.org/wiki/Multiset
    #
    # Outputs guaranteed to only include positive counts.
    #
    # To strip negative and zero counts, add-in an empty counter:
    #       c += Counter()

    def __add__(self, other):
        """Add counts from two counters.

        >>> Counter('abbb') + Counter('bcc')
        Counter({'b': 4, 'c': 2, 'a': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        result = Counter()
        for elem in set(self) | set(other):
            newcount = self[elem] + other[elem]
            if newcount > 0:
                result[elem] = newcount
        return result

    def __sub__(self, other):
        """ Subtract count, but keep only results with positive counts.

        >>> Counter('abbbc') - Counter('bccd')
        Counter({'b': 2, 'a': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        result = Counter()
        for elem in set(self) | set(other):
            newcount = self[elem] - other[elem]
            if newcount > 0:
                result[elem] = newcount
        return result

    def __or__(self, other):
        """Union is the maximum of value in either of the input counters.

        >>> Counter('abbb') | Counter('bcc')
        Counter({'b': 3, 'c': 2, 'a': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        _max = max
        result = Counter()
        for elem in set(self) | set(other):
            newcount = _max(self[elem], other[elem])
            if newcount > 0:
                result[elem] = newcount
        return result

    def __and__(self, other):
        """ Intersection is the minimum of corresponding counts.

        >>> Counter('abbb') & Counter('bcc')
        Counter({'b': 1})

        """
        if not isinstance(other, Counter):
            return NotImplemented
        _min = min
        result = Counter()
        if len(self) < len(other):
            self, other = other, self
        for elem in ifilter(self.__contains__, other):
            newcount = _min(self[elem], other[elem])
            if newcount > 0:
                result[elem] = newcount
        return result

def convert(d):
    "convert a (potentially nested) dict to Storage"
    from types import DictType
    t = type(d)
    if t == DictType:
        for k,v in d.items():
            d[k] = convert(v)
        return Storage(d)
    return d

if __name__ == '__main__':
    import doctest
    print(doctest.testmod())

"""
The Node class and functions for creating trees from Newick strings,
etc.

ivy does not have a Tree class per se, as most functions operate
directly on Node objects.
"""
import os
# from storage import Storage
from copy import copy as _copy
# from matrix import vcv
from . import newick, nexus
# from itertools import izip_longest

## class Tree(object):
##     """
##     A simple Tree class.
##     """
##     def __init__(self, data=None, format="newick", name=None, ttable=None):
##         self.root = None
##         if data:
##             self.root = read(data, format, name, ttable)
##         self.name = name
##         self.ttable = ttable

##     def __getattribute__(self, a):
##         r = object.__getattribute__(self, 'root')
##         try:
##             return object.__getattribute__(r, a)
##         except AttributeError:
##             return object.__getattribute__(self, a)

def traverse(node):
    "recursive preorder iterator based solely on .children attribute"
    yield node
    for child in node.children:
        for descendant in traverse(child):
            yield descendant

class Node(object):
    """
    A basic Node class with attributes and references to child nodes
    ('children', a list) and 'parent'.

    Keyword Args:
        id: ID of the node. If not provided, is set using
          builtin id function
        ni (int): Node index.
        li (int): Leaf index.
        isroot (bool): Is the node a root.
        isleaf (bool): Is the node a leaf.
        label (str): Node label.
        length (float): Branch length from node to parent
        support: RR: Are these bootstrap support values? -CZ
        age (float): Age of the node in time units.
        parent (Node): Parent of the ndoe.
        children (list): List of node objects. Children of node
        nchildren (int): No. of children
        left: RR: Unsure what left and right mean -CZ
        treename: Name of tree
        comment: Comments for tree

    """
    def __init__(self, **kwargs):
        self.id = None
        self.ni = None  # node index
        self.li = None  # leaf index
        self.isroot = False
        self.isleaf = False
        self.label = None
        self.length = None
        self.support = None
        self.age = None
        self.parent = None
        self.children = []
        self.nchildren = 0
        self.left = None
        self.right = None
        self.treename = ""
        self.comment = ""
        self.meta = {}
        ## self.length_comment = ""
        ## self.label_comment = ""
        if kwargs:
            for k, v in kwargs.items():
                setattr(self, k, v)
        if self.id is None:
            self.id = id(self)

    def __copy__(self):
        return self.copy()

    def __repr__(self):
        v = []
        if self.isroot:
            v.append("root")
        elif self.isleaf:
            v.append("leaf")

        if self.label:
            v.append("'%s'" % self.label)

        s = ", ".join(v)

        nid = self.ni if (self.ni is not None) else self.id

        if s:
            s = "Node(%s, %s)" % (nid, s)
        else:
            s = "Node(%s)" % nid
        return s

    def __contains__(self, other):
        """
        For use with `in` keyword

        Args:
            other: Another node or node label.
        Returns:
            bool: Whether or not the other node is a descendant of self
        """
        if other and isinstance(other, str):
            for x in self:
                if other == x.label:
                    return True
            return False
        else:
            assert isinstance(self, Node)
            for x in self.iternodes():
                if other == x:
                    return True
        return False

    def __iter__(self):
        for node in self.iternodes():
            yield node

    def __len__(self):
        """
        Number of nodes descended from self

        Returns:
            int: Number of nodes descended from self (including self)
        """
        i = 0
        for n in self:
            i += 1
        return i

    def __nonzero__(self):
        return True

    def __getitem__(self, x):
        """
        Args:
            x: A Node, Node.id (int) or a Node.label (string)

        Returns:
            Node: Found node(s)

        """
        for n in self:
            if n==x or n.id==x or n.ni == x or (n.label and n.label==x):
                return n
        raise IndexError(str(x))

    def ascii(self, *args, **kwargs):
        """
        Create ascii tree.

        Keyword Args:
            unitlen (float): How long each unit should be rendered as.
              Defaults to 3.
            minwidth (float): Minimum width of the plot. Defaults to 50
            maxwidth (float): Maximum width of the plot. Defaults to None
            scaled (bool): Whether or not the tree is scaled. Defaults to False
            show_internal_labels (bool): Whether or not to show labels
              on internal nodes. Defaults to True.
        Returns:
            str: Ascii tree to be shown with print().
        """
        from . import ascii as _ascii
        return _ascii.render(self, *args, **kwargs)

    def collapse(self, add=False):
        """
        Remove self and collapse children to polytomy

        Args:
            add (bool): Whether or not to add self's length to children's
              length.

        Returns:
            Node: Parent of self

        """
        assert self.parent
        p = self.prune()
        for c in self.children:
            p.add_child(c)
            if add and (c.length is not None):
                c.length += self.length
        self.children = []
        return p

    def copy(self, recurse=False):
        """
        Return a copy of the node, but not copies of children, parent,
        or any attribute that is a Node.

        If `recurse` is True, recursively copy child nodes.

        Args:
            recurse (bool): Whether or not to copy children as well as self.

        Returns:
            Node: A copy of self.

        TODO: test this function.

        RR: This function runs rather slowly -CZ
        """
        newnode = Node()
        for attr, value in self.__dict__.items():
            if (attr not in ("children", "parent") and
                not isinstance(value, Node)):
                setattr(newnode, attr, _copy(value))
            if recurse:
                newnode.children = [
                    child.copy(True) for child in self.children
                    ]
        return newnode

    def leafsets(self, d=None, labels=False):
        """return a mapping of nodes to leaf sets (nodes or labels)"""
        d = d or {}
        if not self.isleaf:
            s = set()
            for child in self.children:
                if child.isleaf:
                    if labels:
                        s.add(child.label)
                    else:
                        s.add(child)
                else:
                    d = child.leafsets(d, labels)
                    s = s | d[child]
            d[self] = frozenset(s)
        return d

    def mrca(self, *nodes):
        """
        Find most recent common ancestor of *nodes*

        Args:
            *nodes (Node): Node objects
        Returns:
            Node: The MRCA of *nodes*
        """
        if len(nodes) == 1:
            nodes = tuple(nodes[0])
        if len(nodes) == 1:
            return nodes[0]
        nodes = set([ self[n] for n in nodes ])
        anc = []
        def f(n):
            seen = set()
            for c in n.children: seen.update(f(c))
            if n in nodes: seen.add(n)
            if seen == nodes and (not anc): anc.append(n)
            return seen
        f(self)
        return anc[0]

    ## def mrca(self, *nodes):
    ##     """
    ##     Find most recent common ancestor of *nodes*
    ##     """
    ##     if len(nodes) == 1:
    ##         nodes = tuple(nodes[0])
    ##     if len(nodes) == 1:
    ##         return nodes[0]
    ##     ## assert len(nodes) > 1, (
    ##     ##     "Need more than one node for mrca(), got %s" % nodes
    ##     ##     )
    ##     def f(x):
    ##         if isinstance(x, Node):
    ##             return x
    ##         elif type(x) in types.StringTypes:
    ##             return self.find(x)
    ##     nodes = map(f, nodes)
    ##     assert all(filter(lambda x: isinstance(x, Node), nodes))

    ##     #v = [ list(n.rootpath()) for n in nodes if n in self ]
    ##     v = [ list(x) for x in izip_longest(*[ reversed(list(n.rootpath()))
    ##                                            for n in nodes if n in self ]) ]
    ##     if len(v) == 1:
    ##         return v[0][0]
    ##     anc = None
    ##     for x in v:
    ##         s = set(x)
    ##         if len(s) == 1: anc = list(s)[0]
    ##         else: break
    ##     return anc

    def ismono(self, *leaves):
        """
        Test if leaf descendants are monophyletic

        Args:
            *leaves (Node): At least two leaf Node objects

        Returns:
            bool: Are the leaf descendants monophyletic?

        RR: Should this function have a check to make sure the input nodes are
        leaves? There is some strange behavior if you input internal nodes -CZ
        """
        if len(leaves) == 1:
            leaves = list(leaves)[0]
        assert len(leaves) > 1, (
            "Need more than one leaf for ismono(), got %s" % leaves
            )
        anc = self.mrca(leaves)
        if anc:
            return bool(len(anc.leaves())==len(leaves))

    def order_subtrees_by_size(self, n2s=None, recurse=False, reverse=False):
        """
        Order interal clades by size

        """
        if n2s is None:
            n2s = clade_sizes(self)
        if not self.isleaf:
            v = [ (n2s[c], c.label or '', c.id, c) for c in self.children ]
            v.sort()
            if reverse:
                v.reverse()
            self.children = [ x[-1] for x in v ]
            if recurse:
                for c in self.children:
                    c.order_subtrees_by_size(n2s, recurse=True, reverse=reverse)

    def ladderize(self, reverse=False):
        self.order_subtrees_by_size(recurse=True, reverse=reverse)
        return self

    def add_child(self, child):
        """
        Add child as child of self

        Args:
            child (Node): A node object

        """
        assert child not in self.children
        self.children.append(child)
        child.parent = self
        child.isroot = False
        self.nchildren += 1

    def bisect_branch(self):
        """
        Add new node as parent to self in the middle of branch to parent.

        Returns:
            Node: A new node.

        """
        assert self.parent
        parent = self.prune()
        n = Node()
        if self.length:
            n.length = self.length/2.0
            self.length /= 2.0
        parent.add_child(n)
        n.add_child(self)
        return n

    def remove_child(self, child):
        """
        Remove child.

        Args:
            child (Node): A node object that is a child of self

        """
        assert child in self.children
        self.children.remove(child)
        child.parent = None
        self.nchildren -= 1
        if not self.children:
            self.isleaf = True

    def labeled(self):
        """
        Return a list of all descendant nodes that are labeled

        Returns:
            list: All descendants of self that are labeled (including self)
        """
        return [ n for n in self if n.label ]

    def leaves(self, f=None):
        """
        Return a list of leaves. Can be filtered with f.

        Args:
            f (function): A function that evaluates to True if called with desired
              node as the first input

        Returns:
            list: A list of leaves that are true for f (if f is given)

        """
        if f: return [ n for n in self if (n.isleaf and f(n)) ]
        return [ n for n in self if n.isleaf ]

    def internals(self, f=None):
        """
        Return a list nodes that have children (internal nodes)

        Args:
            f (function): A function that evaluates to true if called with desired
              node as the first input

        Returns:
            list: A list of internal nodes that are true for f (if f is given)

        """
        if f: return [ n for n in self if (n.children and f(n)) ]
        return [ n for n in self if n.children ]

    def clades(self):
        """
        Get internal nodes descended from self

        Returns:
            list: A list of internal nodes descended from (and not including) self.

        """
        return [ n for n in self if (n is not self) and not n.isleaf ]

    def iternodes(self, f=None):
        """
        Return a generator of nodes descendant from self - including self

        Args:
            f (function): A function that evaluates to true if called with
            desired node as the first input

        Yields:
            Node: Nodes descended from self (including self) in
              preorder sequence

        """
        if (f and f(self)) or (not f):
            yield self
        for child in self.children:
            for n in child.iternodes(f):
                yield n

    def iterleaves(self):
        """
        Yield leaves descendant from self
        """
        return self.iternodes(lambda x:x.isleaf)

    def preiter(self, f=None):
        """
        Yield nodes in preorder sequence
        """
        for n in self.iternodes(f=f):
            yield n

    def postiter(self, f=None):
        """
        Yield nodes in postorder sequence
        """
        if not self.isleaf:
            for child in self.children:
                for n in child.postiter():
                    if (f and f(n)) or (not f):
                        yield n
        if (f and f(self)) or (not f):
            yield self

    def descendants(self, order="pre", v=None, f=None):
        """
        Return a list of nodes descendant from self - but _not_
        including self!

        Args:
            order (str): Indicates wether to return nodes in preorder or
              postorder sequence. Optional, defaults to "pre"
            f (function): filtering function that evaluates to True if desired
              node is called as the first parameter.

        Returns:
            list: A list of nodes descended from self not including self.

        """
        v = v or []
        for child in self.children:
            if (f and f(child)) or (not f):
                if order == "pre":
                    v.append(child)
                else:
                    v.insert(0, child)
            if child.children:
                child.descendants(order, v, f)
        return v

    def get(self, f, *args, **kwargs):
        """
        Return the first node found by node.find()

        Args:
            f (function): A function that evaluates to True if desired
              node is called as the first parameter.
        Returns:
            Node: The first node found by node.find()

        """
        v = self.find(f, *args, **kwargs)
        try:
            return v.next()
        except StopIteration:
            return None

    def grep(self, s, ignorecase=True):
        """
        Find nodes by regular-expression search of labels

        Args:
            s (str): String to search.
            ignorecase (bool): Indicates to ignore case. Defaults to true.

        Returns:
            lsit: A list of node objects whose labels were matched by s.

        """
        import re
        if ignorecase:
            pattern = re.compile(s, re.IGNORECASE)
        else:
            pattern = re.compile(s)

        search = pattern.search
        return [ x for x in self if x.label and search(x.label) ]

    def lgrep(self, s, ignorecase=True):
        """
        Find leaves by regular-expression search of labels

        Args:
            s (str): String to search.
            ignorecase (bool): Indicates to ignore case. Defaults to true.

        Returns:
            lsit: A list of node objects whose labels were matched by s.

        """
        return [ x for x in self.grep(s, ignorecase=ignorecase) if x.isleaf ]

    def bgrep(self, s, ignorecase=True):
        """
        Find branches (internal nodes) by regular-expression search of
        labels

        Args:
            s (str): String to search.
            ignorecase (bool): Indicates to ignore case. Defaults to true.

        Returns:
            lsit: A list of node objects whose labels were matched by s.

        """
        return [ x for x in self.grep(s, ignorecase=ignorecase) if
               (not x.isleaf) ]

    def find(self, f, *args, **kwargs):
        """
        Find descendant nodes.

        Args:
            f: Function or a string.  If a string, it is converted to a
              function for finding *f* as a substring in node labels.
              Otherwise, *f* should evaluate to True if called with a desired
              node as the first parameter.

        Yields:
            Node: Found nodes in preorder sequence.

        """
        if not f:
            return
        if isinstance(f, str):
            func = lambda x: (f or None) in (x.label or "")
        else:
            func = f
        for n in self.iternodes():
            if func(n, *args, **kwargs):
                yield n

    def findall(self, f, *args, **kwargs):
        """Return a list of found nodes."""
        return list(self.find(f, *args, **kwargs))

    def prune(self):
        """
        Remove self if self is not root.

        Returns:
            Node: Parent of self. If parent had only two children,
              parent is now a 'knee' and can be removed with excise.

        """
        p = self.parent
        if p:
            p.remove_child(self)
        return p

    def excise(self):
        """
        For 'knees': remove self from between parent and single child
        """
        assert self.parent
        assert len(self.children)==1
        p = self.parent
        c = self.children[0]
        if c.length is not None and self.length is not None:
            c.length += self.length
        c.prune()
        self.prune()
        p.add_child(c)
        return p

    def graft(self, node):
        """
        Add node as sister to self.
        """
        parent = self.parent
        parent.remove_child(self)
        n = Node()
        n.add_child(self)
        n.add_child(node)
        parent.add_child(n)

    ## def leaf_distances(self, store=None, measure="length"):
    ##     """
    ##     for each internal node, calculate the distance to each leaf,
    ##     measured in branch length or internodes
    ##     """
    ##     if store is None:
    ##         store = {}
    ##     leaf2len = {}
    ##     if self.children:
    ##         for child in self.children:
    ##             if measure == "length":
    ##                 dist = child.length
    ##             elif measure == "nodes":
    ##                 dist = 1
    ##             child.leaf_distances(store, measure)
    ##             if child.isleaf:
    ##                 leaf2len[child] = dist
    ##             else:
    ##                 for k, v in store[child].items():
    ##                     leaf2len[k] = v + dist
    ##     else:
    ##         leaf2len[self] = {self: 0}
    ##     store[self] = leaf2len
    ##     return store

    def leaf_distances(self, measure="length"):
        """
        RR: I don't quite understand the structure of the output. Also,
            I can't figure out what "measure" does.-CZ
        """
        from collections import defaultdict
        store = defaultdict(lambda:defaultdict(lambda:0))
        nodes = [ x for x in self if x.children ]
        for lf in self.leaves():
            x = lf.length
            for n in lf.rootpath(self):
                store[n][lf] = x
                x += (n.length or 0)
        return store

    def rootpath(self, end=None, stop=None):
        """
        Iterate over parent nodes toward the root, or node *end* if
        encountered.

        Args:
            end (Node): A Node object to iterate to (instead of iterating
              towards root). Optional, defaults to None
            stop (function): A function that returns True if desired node is called
              as the first parameter. Optional, defaults to None

        Yields:
            Node: Nodes in path to root (or end).

        """
        n = self.parent
        while 1:
            if n is None: raise StopIteration
            yield n
            if n.isroot or (end and n == end) or (stop and stop(n)):
                raise StopIteration
            n = n.parent

    def rootpath_length(self, end=None):
        """
        Get length from self to root(if end is None) or length
        from self to an ancestor node (if end is an ancestor to self)

        Args:
            end (Node): A node object

        Returns:
            float: The length from self to root/end

        """
        n = self
        x = 0.0
        while n.parent:
            x += n.length
            if n.parent == end:
                break
            n = n.parent
        return x
        ## f = lambda x:x.parent==end
        ## v = [self.length]+[ x.length for x in self.rootpath(stop=f)
        ##                     if x.parent ]
        ## assert None not in v
        ## return sum(v)

    def max_tippath(self, first=True):
        """
        Get the maximum length from self to a leaf node
        """
        v = 0
        if self.children:
            v = max([ c.max_tippath(False) for c in self.children ])
        if not first:
            if self.length is None: v += 1
            else: v += self.length
        return v

    def subtree_mapping(self, labels, clean=False):
        """
        Find the set of nodes in 'labels', and create a new tree
        representing the subtree connecting them.  Nodes are assumed
        to be non-nested.

        Returns:
            dict: a mapping of old nodes to new nodes and vice versa.

        TODO: test this, high bug probability
        """
        d = {}
        oldtips = [ x for x in self.leaves() if x.label in labels ]
        for tip in oldtips:
            path = list(tip.rootpath())
            for node in path:
                if node not in d:
                    newnode = Node()
                    newnode.isleaf = node.isleaf
                    newnode.length = node.length
                    newnode.label = node.label
                    d[node] = newnode
                    d[newnode] = node
                else:
                    newnode = d[node]

                for child in node.children:
                    if child in d:
                        newchild = d[child]
                        if newchild not in newnode.children:
                            newnode.add_child(newchild)
        d["oldroot"] = self
        d["newroot"] = d[self]
        if clean:
            n = d["newroot"]
            while 1:
                if n.nchildren == 1:
                    oldnode = d[n]
                    del d[oldnode]; del d[n]
                    child = n.children[0]
                    child.parent = None
                    child.isroot = True
                    d["newroot"] = child
                    d["oldroot"] = d[child]
                    n = child
                else:
                    break

            for tip in oldtips:
                newnode = d[tip]
                while 1:
                    newnode = newnode.parent
                    oldnode = d[newnode]
                    if newnode.nchildren == 1:
                        child = newnode.children[0]
                        if newnode.length:
                            child.length += newnode.length
                        newnode.remove_child(child)
                        if newnode.parent:
                            parent = newnode.parent
                            parent.remove_child(newnode)
                            parent.add_child(child)
                        del d[oldnode]; del d[newnode]
                    if not newnode.parent:
                        break

        return d

    def reroot_orig(self, newroot):
        assert newroot in self
        self.isroot = False
        newroot.isroot = True
        v = []
        n = newroot
        while 1:
            v.append(n)
            if not n.parent: break
            n = n.parent
        v.reverse()
        for i, cp in enumerate(v[:-1]):
            node = v[i+1]
            # node is current node; cp is current parent
            cp.remove_child(node)
            node.add_child(cp)
            cp.length = node.length
        return newroot

    def reroot(self, newroot):
        """
        RR: I can't get this to work properly -CZ
        """
        newroot = self[newroot]
        assert newroot in self
        self.isroot = False
        n = newroot
        v = list(n.rootpath())
        v.reverse()
        for node in (v+[n])[1:]:
            # node is current node; cp is current parent
            cp = node.parent
            cp.remove_child(node)
            node.add_child(cp)
            cp.length = node.length
            cp.label = node.label
        newroot.isroot = True
        return newroot

    def makeroot(self, shift_labels=False):
        """
        shift_labels: flag to shift internal parent-child node labels
        when internode polarity changes; suitable e.g. if internal node
        labels indicate unrooted bipartition support
        """
        v = list(self.rootpath())
        v[-1].isroot = False
        v.reverse()
        for node in v[1:] + [self]:
            # node is current node; cp is current parent
            cp = node.parent
            cp.remove_child(node)
            node.add_child(cp)
            cp.length = node.length
            if shift_labels:
                cp.label = node.label
        self.isroot = True
        return self

    def write(self, outfile=None, format="newick", length_fmt=":%g", end=True,
              clobber=False):
        if format=="newick":
            s = write_newick(self, outfile, length_fmt, True, clobber)
            if not outfile:
                return s

reroot = Node.reroot

def index(node, n=0, d=0):
    """
    recursively attach 'next', 'back', (and 'left', 'right'), 'ni',
    'ii', 'pi', and 'node_depth' attributes to nodes
    """
    node.next = node.left = n
    if not node.parent:
        node.node_depth = d
    else:
        node.node_depth = node.parent.node_depth + 1
    n += 1
    for i, c in enumerate(node.children):
        if i > 0:
            n = node.children[i-1].back + 1
        index(c, n)

    if node.children:
        node.back = node.right = node.children[-1].back + 1
    else:
        node.back = node.right = n
    return node.back

def remove_singletons(root, add=True):
    "Remove descendant nodes that are the sole child of their parent"
    for leaf in root.leaves():
        for n in leaf.rootpath():
            if n.parent and len(n.parent.children)==1:
                n.collapse(add)

def cls(root):
    """
    Get clade sizes of whole tree
    Args:
        * root: A root node

    Returns:
        * A dict mapping nodes to clade sizes

    """
    results = {}
    for node in root.postiter():
        if node.isleaf:
            results[node] = 1
        else:
            results[node] = sum(results[child] for child in node.children)
    return results

def clade_sizes(node, results={}):
    """Map node and descendants to number of descendant tips"""
    size = int(node.isleaf)
    if not node.isleaf:
        for child in node.children:
            clade_sizes(child, results)
            size += results[child]
    results[node] = size
    return results

def write(node, outfile=None, format="newick", length_fmt=":%g",
          clobber=False):
    if format=="newick" or (isinstance(outfile, str) and
                            outfile.endswith(".newick") or
                            outfile.endswith(".new")):
        s = write_newick(node, outfile, length_fmt, True, clobber)
        if not outfile:
            return s

def write_newick(node, outfile=None, length_fmt=":%g", end=False,
                 clobber=False):
    if not node.isleaf:
        node_str = "(%s)%s" % \
                   (",".join([ write_newick(child, outfile, length_fmt,
                                            False, clobber)
                               for child in node.children ]),
                    (node.label or "")
                    )
    else:
        node_str = "%s" % node.label

    if node.length is not None:
        length_str = length_fmt % node.length
    else:
        length_str = ""

    semicolon = ""
    if end:
        semicolon = ";"
    s = "%s%s%s" % (node_str, length_str, semicolon)
    if end and outfile:
        flag = False
        if isinstance(outfile, str):
            if not clobber:
                assert not os.path.isfile(outfile), "File '%s' exists! (Set clobber=True to overwrite)" % outfile
            flag = True
            outfile = open(outfile, "w")
        outfile.write(s)
        if flag:
            outfile.close()
    return s

def read(data, format=None, treename=None, ttable=None):
    """
    Read a single tree from *data*, which can be a Newick string, a
    file name, or a file-like object with `tell` and 'read`
    methods. *treename* is an optional string that will be attached to
    all created nodes.

    Args:
        data: A file or file-like object or newick string

    Returns:
        Node: The root node.
    """

    def strip(s):
        fname = os.path.split(s)[-1]
        head, tail = os.path.splitext(fname)
        tail = tail.lower()
        if tail in (".nwk", ".tre", ".tree", ".newick", ".nex"):
            return head
        else:
            return fname

    if (not format):
        if isinstance(data, str) and os.path.isfile(data):
            s = data.lower()
            for tail in ".nex", ".nexus", ".tre":
                if s.endswith(tail):
                    format="nexus"
                    break

    if (not format):
        format = "newick"

    if format == "newick":
        if isinstance(data, str):
            if os.path.isfile(data):
                treename = strip(data)
                return newick.parse(open(data), treename=treename,
                                    ttable=ttable)
            else:
                return newick.parse(data, ttable=ttable)

        elif (hasattr(data, "tell") and hasattr(data, "read")):
            treename = strip(getattr(data, "name", None))
            return newick.parse(data, treename=treename, ttable=ttable)
    elif format == "nexus-dendropy":
        import dendropy
        if isinstance(data, str):
            if os.path.isfile(data):
                treename = strip(data)
                return newick.parse(
                    str(dendropy.Tree.get_from_path(data, "nexus")),
                    treename=treename
                    )
            else:
                return newick.parse(
                    str(dendropy.Tree.get_from_string(data, "nexus"))
                    )

        elif (hasattr(data, "tell") and hasattr(data, "read")):
            treename = strip(getattr(data, "name", None))
            return newick.parse(
                str(dendropy.Tree.get_from_stream(data, "nexus")),
                treename=treename
                )
        else:
            pass

    elif format == "nexus":
        if isinstance(data, str):
            if os.path.isfile(data):
                with open(data) as infile:
                    nexiter = nexus.iter_trees(infile)
                    rec = next(nexiter)
                    if rec:
                        return rec.parse()
            else:
                nexiter = nexus.iter_trees(StringIO(data))
        else:
            nexiter = nexus.iter_trees(data)
        rec = next(nexiter)
        if rec:
            return rec.parse()
    else:
        # implement other tree formats here (nexus, nexml etc.)
        raise IOError("format '%s' not implemented yet" % format)

    raise IOError("unable to read tree from '%s'" % data)

def readmany(data, format="newick"):
    """Iterate over trees from a source."""
    if isinstance(data, str):
        if os.path.isfile(data):
            data = open(data)
        else:
            data = StringIO(data)

    if format == "newick":
        for line in data:
            yield newick.parse(line)
    elif format == "nexus":
        for rec in newick.nexus_iter(data):
            yield rec.parse()
    else:
        raise Exception("format '%s' not recognized" % format)
    data.close()

## def randomly_resolve(n):
##     assert len(n.children)>2

## def leaf_mrcas(root):
##     from itertools import product, izip, tee
##     from collections import OrderedDict
##     from numpy import empty
##     mrca = OrderedDict()
##     def pairwise(iterable, tee=tee, izip=izip):
##         a, b = tee(iterable)
##         next(b, None)
##         return izip(a, b)
##     def f(n):
##         if n.isleaf:
##             od = OrderedDict(); od[n] = n.length
##             return od
##         d = [ f(c) for c in n.children ]
##         for i, j in pairwise(xrange(len(d))):
##             di = d[i]; dj =d[j]
##             for ni, niv in di.iteritems():
##                 for nj, njv in dj.iteritems():
##                     mrca[(ni,nj)] = n
##             d[j].update(di)
##         return d[j]
##     f(root)
##     return mrca

def C(leaves, internals):
    from scipy.sparse import lil_matrix
    m = lil_matrix((len(internals), len(leaves)))
    for lf in leaves:
        v = lf.length if lf.length is not None else 1
        for n in lf.rootpath():
            m[n.ii,lf.li] = v
            v += n.length if n.length is not None else 1
    return m.tocsc()

"""
Functions to get trees and character data from treebase
"""

try:
    from urllib2 import urlopen
except ImportError:
    from urllib.request import urlopen
from lxml import etree
from . import storage
import sys, re

# "http://purl.org/phylo/treebase/phylows/study/TB2:S11152"

TREEBASE_WEBSERVICE = "http://purl.org/phylo/treebase/phylows"
NEXML_NAMESPACE = "http://www.nexml.org/2009"
NEXML = "{%s}" % NEXML_NAMESPACE
UNIPROT = "http://purl.uniprot.org/taxonomy/"
NAMEBANK = ("http://www.ubio.org/authority/metadata.php?"
            "lsid=urn:lsid:ubio.org:namebank:")

ROW_SEGMENTS = ("http://treebase.org/treebase-web/search/study/"
                "rowSegmentsTSV.html?matrixid=")

META_DATATYPE = {
    "xsd:long": int,
    "xsd:integer": int,
    "xsd:string": str
    }

AMBIG_RE = re.compile(r'([{][a-zA-Z]+[}])')

def fetch_study(study_id, format="nexml"):
    """
    Get a study from treebase in one of various formats

    Args:
        study_id (str): The id of the study
        format (str):  One of ["rdf", "html", "nexml", "nexus"]
    Returns:
        Str representing a nexus file (if format = "nexus")

        OR

        An lxml etree object
    """
    try: study_id = "S%s" % int(study_id)
    except ValueError: pass

    # format is one of ["rdf", "html", "nexml", "nexus"]
    url = "%s/study/TB2:%s?format=%s" % (TREEBASE_WEBSERVICE, study_id, format)
    if format=="nexus":
        return urlopen(url).read()
    else:
        return etree.parse(url)

def parse_chars(e, otus):
    v = []
    for chars in e.findall(NEXML+"characters"):
        c = storage.Storage(chars.attrib)
        c.states = parse_states(chars)
        c.meta = storage.Storage()
        for meta in chars.findall(NEXML+"meta"):
            a = meta.attrib
            if a.get("content"):
                value = META_DATATYPE[a["datatype"]](a["content"])
                c.meta[a["property"]] = value
        c.matrices = []
        for matrix in chars.findall(NEXML+"matrix"):
            m = storage.Storage()
            m.rows = []
            for row in matrix.findall(NEXML+"row"):
                r = storage.Storage(row.attrib)
                r.otu = otus[r.otu]
                s = row.findall(NEXML+"seq")[0].text
                substrs = []
                for ss in AMBIG_RE.split(s):
                    if ss.startswith("{"):
                        key = frozenset(ss[1:-1])
                        val = c.states.states2symb.get(key)
                        if key and not val:
                            sys.stderr.write("missing ambig symbol for %s\n" %
                                             "".join(sorted(key)))
                        ss = val or "?"
                    substrs.append(ss)
                s = "".join(substrs)
                r.seq = s
                m.rows.append(r)
            c.matrices.append(m)
        v.append(c)
    return v

def parse_trees(e, otus):
    """
    Get trees from an etree object

    Args:
        e: A nexml document parsed by etree
        otus: OTUs returned by parse_otus
    Returns:
        list: A list of ivy Storage objects each
          containing every node of a tree.
    """
    from tree import Node
    v = []
    for tb in e.findall(NEXML+"trees"):
        for te in tb.findall(NEXML+"tree"):
            t = storage.Storage()
            t.attrib = storage.Storage(te.attrib)
            t.nodes = {}
            for n in te.findall(NEXML+"node"):
                node = Node()
                if n.attrib.get("otu"):
                    node.isleaf = True
                    node.otu = otus[n.attrib["otu"]]
                    node.label = node.otu.label
                t.nodes[n.attrib["id"]] = node
            for edge in te.findall(NEXML+"edge"):
                d = edge.attrib
                n = t.nodes[d["target"]]
                p = t.nodes[d["source"]]
                length = d.get("length")
                if length:
                    n.length = float(length)
                p.add_child(n)
            r = [ n for n in t.nodes.values() if not n.parent ]
            assert len(r)==1
            r = r[0]
            r.isroot = True
            for i, n in enumerate(r): n.id = i+1
            t.root = r
            v.append(t)
    return v

def parse_otus(e):
    """
    Get OTUs from an etree object

    Args:
        e: A nexml document parsed by etree
    Returns:
        dict: A dict mapping keys to OTUs contained in ivy Storage objects
    """
    v = {}
    for otus in e.findall(NEXML+"otus"):
        for x in otus.findall(NEXML+"otu"):
            otu = storage.Storage()
            otu.id = x.attrib["id"]
            otu.label = x.attrib["label"]
            for meta in x.iterchildren():
                d = meta.attrib
                p = d.get("property")
                if p and p == "tb:identifier.taxon":
                    otu.tb_taxid = d["content"]
                elif p and p == "tb:identifier.taxonVariant":
                    otu.tb_taxid_variant = d["content"]
                h = d.get("href")
                if h and h.startswith(NAMEBANK):
                    otu.namebank_id = int(h.replace(NAMEBANK, ""))
                elif h and h.startswith(UNIPROT):
                    otu.ncbi_taxid = int(h.replace(UNIPROT, ""))
            v[otu.id] = otu
    return v

def parse_nexml(doc):
    """
    Parse an etree ElementTree

    Args:
        doc: An etree ElementTree or a file that can be parsed into
          an etree ElementTree with etree.parse
    Returns:
        An ivy Storage object containing all the information from the
          nexml file: Characters, metadata, OTUs, and trees.
    """
    if not isinstance(doc, (etree._ElementTree, etree._Element)):
        doc = etree.parse(doc)
    meta = {}
    for child in doc.findall(NEXML+"meta"):
        if "content" in child.attrib:
            d = child.attrib
            key = d["property"]
            val = META_DATATYPE[d["datatype"]](d["content"])
            if (key in meta) and val:
                if isinstance(meta[key], list):
                    meta[key].append(val)
                else:
                    meta[key] = [meta[key], val]
            else:
                meta[key] = val

    otus = parse_otus(doc)

    return Storage(meta = meta,
                   otus = otus,
                   chars = parse_chars(doc, otus),
                   trees = parse_trees(doc, otus))

def parse_states(e):
    """e is a characters element"""
    f = e.findall(NEXML+"format")[0]
    sts = f.findall(NEXML+"states")[0]
    states2symb = {}
    symb2states = {}
    id2symb = {}
    for child in sts.iterchildren():
        t = child.tag
        if t == NEXML+"state":
            k = child.attrib["id"]
            v = child.attrib["symbol"]
            id2symb[k] = v
            states2symb[v] = v
            symb2states[v] = v
        elif t == NEXML+"uncertain_state_set":
            k = child.attrib["id"]
            v = child.attrib["symbol"]
            id2symb[k] = v
            memberstates = []
            for memb in child.findall(NEXML+"member"):
                sid = memb.attrib["state"]
                symb = id2symb[sid]
                for x in symb2states[symb]: memberstates.append(x)
            memberstates = frozenset(memberstates)
            symb2states[v] = memberstates
            states2symb[memberstates] = v
    return Storage(states2symb=states2symb,
                   symb2states=symb2states,
                   id2symb=id2symb)

def parse_charsets(study_id):
    from cStringIO import StringIO
    nx = StringIO(fetch_study(study_id, 'nexus'))
    d = {}
    for line in nx.readlines():
        if line.strip().startswith("CHARSET "):
            v = line.strip().split()
            label = v[1]
            first, last = map(int, line.split()[-1][:-1].split("-"))
            d[label] = (first-1, last-1)
    return d

if __name__ == "__main__":
    import sys
    from pprint import pprint
    e = fetch_study('S11152', 'nexus')
    #print e
    #e.write(sys.stdout, pretty_print=True)

    ## e = etree.parse('/tmp/tmp.xml')
    ## x = parse_nexml(e)
    ## pprint(x)

"""
interactive viewers for trees, etc. using matplotlib
"""
from collections import defaultdict
from .. import align, sequtil
import matplotlib, numpy
import matplotlib.pyplot as pyplot
from matplotlib.figure import SubplotParams
from matplotlib.axes import Axes, subplot_class_factory
from matplotlib.patches import Rectangle
from matplotlib.collections import PatchCollection
from matplotlib.widgets import RectangleSelector
from matplotlib import colors as mpl_colors
from matplotlib.ticker import MaxNLocator, FuncFormatter, NullLocator
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from Bio.Align import MultipleSeqAlignment

matplotlib.rcParams['path.simplify'] = False

class UpdatingRect(Rectangle):
    def __call__(self, p):
        self.set_bounds(*p.viewLim.bounds)
        p.figure.canvas.draw_idle()

class AlignmentFigure:
    def __init__(self, aln, name=None, div=0.25, overview=True):
        if isinstance(aln, MultipleSeqAlignment):
            self.aln = aln
        else:
            self.aln = align.read(aln)
        self.name = name
        self.div_value = div
        pars = SubplotParams(
            left=0.2, right=1, bottom=0.05, top=1, wspace=0.01
            )
        fig = pyplot.figure(subplotpars=pars, facecolor="white")
        self.figure = fig
        self.initialize_subplots(overview)
        self.show()
        self.connect_events()

    def initialize_subplots(self, overview=False):
        ## p = AlignmentPlot(self.figure, 212, aln=self.aln)
        p = AlignmentPlot(self.figure, 111, aln=self.aln, app=self)
        self.detail = self.figure.add_subplot(p)
        self.detail.plot_aln()
        if overview:
            self.overview = inset_axes(
                self.detail, width="30%", height="20%", loc=1
                )
            self.overview.xaxis.set_major_locator(NullLocator())
            self.overview.yaxis.set_major_locator(NullLocator())
            self.overview.imshow(
                self.detail.array, interpolation='nearest', aspect='auto',
                origin='lower'
                )
            rect = UpdatingRect(
                [0,0], 0, 0, facecolor='black', edgecolor='cyan', alpha=0.5
                )
            self.overview.zoomrect = rect
            rect.target = self.detail
            self.detail.callbacks.connect('xlim_changed', rect)
            self.detail.callbacks.connect('ylim_changed', rect)
            self.overview.add_patch(rect)
            rect(self.overview)

        else:
            self.overview = None

    def show(self):
        self.figure.show()

    def connect_events(self):
        mpl_connect = self.figure.canvas.mpl_connect
        mpl_connect("button_press_event", self.onclick)
        mpl_connect("button_release_event", self.onbuttonrelease)
        mpl_connect("scroll_event", self.onscroll)
        mpl_connect("pick_event", self.onpick)
        mpl_connect("motion_notify_event", self.ondrag)
        mpl_connect("key_press_event", self.onkeypress)
        mpl_connect("axes_enter_event", self.axes_enter)
        mpl_connect("axes_leave_event", self.axes_leave)

    @staticmethod
    def axes_enter(e):
        ax = e.inaxes
        ax._active = True

    @staticmethod
    def axes_leave(e):
        ax = e.inaxes
        ax._active = False

    @staticmethod
    def onselect(estart, estop):
        b = estart.button
        ## print b, estart.key

    @staticmethod
    def onkeypress(e):
        ax = e.inaxes
        k = e.key
        if ax and k:
            if k == 't':
                ax.home()
            elif k == "down":
                ax.scroll(0, -0.1)
            elif k == "up":
                ax.scroll(0, 0.1)
            elif k == "left":
                ax.scroll(-0.1, 0)
            elif k == "right":
                ax.scroll(0.1, 0)
            elif k in '=+':
                ax.zoom(0.1,0.1)
            elif k == '-':
                ax.zoom(-0.1,-0.1)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def ondrag(e):
        ax = e.inaxes
        button = e.button
        if ax and button == 2:
            if not ax.pan_start:
                ax.pan_start = (e.xdata, e.ydata)
                return
            x, y = ax.pan_start
            xdelta = x - e.xdata
            ydelta = y - e.ydata
            x0, x1 = ax.get_xlim()
            xspan = x1-x0
            y0, y1 = ax.get_ylim()
            yspan = y1 - y0
            midx = (x1+x0)*0.5
            midy = (y1+y0)*0.5
            ax.set_xlim(midx+xdelta-xspan*0.5, midx+xdelta+xspan*0.5)
            ax.set_ylim(midy+ydelta-yspan*0.5, midy+ydelta+yspan*0.5)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onbuttonrelease(e):
        ax = e.inaxes
        button = e.button
        if button == 2:
            ## print "pan end"
            ax.pan_start = None
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onpick(e):
        ax = e.mouseevent.inaxes
        if ax:
            ax.picked(e)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onscroll(e):
        ax = e.inaxes
        if ax:
            b = e.button
            ## print b
            k = e.key
            if k == None and b =="up":
                ax.zoom(0.1,0.1)
            elif k == None and b =="down":
                ax.zoom(-0.1,-0.1)
            elif k == "shift" and b == "up":
                ax.zoom_cxy(0.1, 0, e.xdata, e.ydata)
            elif k == "shift" and b == "down":
                ax.zoom_cxy(-0.1, 0, e.xdata, e.ydata)
            elif k == "control" and b == "up":
                ax.zoom_cxy(0, 0.1, e.xdata, e.ydata)
            elif k == "control" and b == "down":
                ax.zoom_cxy(0, -0.1, e.xdata, e.ydata)
            elif k == "d" and b == "up":
                ax.scroll(0, 0.1)
            elif (k == "d" and b == "down"):
                ax.scroll(0, -0.1)
            elif k == "c" and b == "up":
                ax.scroll(-0.1, 0)
            elif k == "c" and b == "down":
                ax.scroll(0.1, 0)
            ax.figure.canvas.draw_idle()

    @staticmethod
    def onclick(e):
        ax = e.inaxes
        if (ax and e.button==1 and hasattr(ax, "zoomrect") and ax.zoomrect):
            # overview clicked; reposition zoomrect
            r = ax.zoomrect
            x = e.xdata
            y = e.ydata
            arr = ax.transData.inverted().transform(r.get_extents())
            xoff = (arr[1][0]-arr[0][0])*0.5
            yoff = (arr[1][1]-arr[0][1])*0.5
            r.target.set_xlim(x-xoff,x+xoff)
            r.target.set_ylim(y-yoff,y+yoff)
            r(r.target)
            ax.figure.canvas.draw_idle()

        elif ax and e.button==2:
            ## print "pan start", (e.xdata, e.ydata)
            ax.pan_start = (e.xdata, e.ydata)
            ax.figure.canvas.draw_idle()

        elif ax and hasattr(ax, "aln") and ax.aln:
            x = int(e.xdata+0.5); y = int(e.ydata+0.5)
            aln = ax.aln
            if (0 <= x <= ax.nchar) and (0 <= y <= ax.ntax):
                seq = aln[y]; char = seq[x]
                if char not in '-?':
                    v = sequtil.gapidx(seq)
                    i = (v[1]==x).nonzero()[0][0]
                    print("%s: row %s, site %s: '%s', seqpos %s"
                           % (seq.id, y, x, char, i))
                else:
                    print("%s: row %s, site %s: '%s'" % (seq.id, y, x, char))

    def zoom(self, factor=0.1):
        "Zoom both axes by *factor* (relative display size)."
        self.detail.zoom(factor, factor)
        self.figure.canvas.draw_idle()

    def __get_selection(self):
        return self.detail.extract_selected()
    selected = property(__get_selection)

class Alignment(Axes):
    """
    matplotlib.axes.Axes subclass for rendering sequence alignments.
    """
    def __init__(self, fig, rect, *args, **kwargs):
        self.aln = kwargs.pop("aln")
        nrows = len(self.aln)
        ncols = self.aln.get_alignment_length()
        self.alnidx = numpy.arange(ncols)
        self.app = kwargs.pop("app", None)
        self.showy = kwargs.pop('showy', True)
        Axes.__init__(self, fig, rect, *args, **kwargs)
        rgb = mpl_colors.colorConverter.to_rgb
        gray = rgb('gray')
        d = defaultdict(lambda:gray)
        d["A"] = rgb("red")
        d["a"] = rgb("red")
        d["C"] = rgb("blue")
        d["c"] = rgb("blue")
        d["G"] = rgb("green")
        d["g"] = rgb("green")
        d["T"] = rgb("yellow")
        d["t"] = rgb("yellow")
        self.cmap = d
        self.selector = RectangleSelector(
            self, self.select_rectangle, useblit=True
            )
        def f(e):
            if e.button != 1: return True
            else: return RectangleSelector.ignore(self.selector, e)
        self.selector.ignore = f
        self.selected_rectangle = Rectangle(
            [0,0],0,0, facecolor='white', edgecolor='cyan', alpha=0.3
            )
        self.add_patch(self.selected_rectangle)
        self.highlight_find_collection = None

    def plot_aln(self):
        cmap = self.cmap
        self.ntax = len(self.aln); self.nchar = self.aln.get_alignment_length()
        a = numpy.array([ [ cmap[base] for base in x.seq ]
                          for x in self.aln ])
        self.array = a
        self.imshow(a, interpolation='nearest', aspect='auto', origin='lower')
        y = [ i+0.5 for i in xrange(self.ntax) ]
        labels = [ x.id for x in self.aln ]
        ## locator.bin_boundaries(1,ntax)
        ## locator.view_limits(1,ntax)
        if self.showy:
            locator = MaxNLocator(nbins=50, integer=True)
            self.yaxis.set_major_locator(locator)
            def fmt(x, pos=None):
                if x<0: return ""
                try: return labels[int(round(x))]
                except: pass
                return ""
            self.yaxis.set_major_formatter(FuncFormatter(fmt))
        else:
            self.yaxis.set_major_locator(NullLocator())

        return self

    def select_rectangle(self, e0, e1):
        x0, x1 = map(int, sorted((e0.xdata+0.5, e1.xdata+0.5)))
        y0, y1 = map(int, sorted((e0.ydata+0.5, e1.ydata+0.5)))
        self.selected_chars = (x0, x1)
        self.selected_taxa = (y0, y1)
        self.selected_rectangle.set_bounds(x0-0.5,y0-0.5,x1-x0+1,y1-y0+1)
        self.app.figure.canvas.draw_idle()

    def highlight_find(self, substr):
        if not substr:
            if self.highlight_find_collection:
                self.highlight_find_collection.remove()
                self.highlight_find_collection = None
            return

        N = len(substr)
        v = []
        for y, x in align.find(self.aln, substr):
            r = Rectangle(
                [x-0.5,y-0.5], N, 1,
                facecolor='cyan', edgecolor='cyan', alpha=0.7
                )
            v.append(r)
        if self.highlight_find_collection:
            self.highlight_find_collection.remove()
        c = PatchCollection(v, True)
        self.highlight_find_collection = self.add_collection(c)
        self.app.figure.canvas.draw_idle()

    def extract_selected(self):
        r0, r1 = self.selected_taxa
        c0, c1 = self.selected_chars
        return self.aln[r0:r1+1,c0:c1+1]

    def zoom_cxy(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view, with a fixed data point (cx, cy)
        """
        transform = self.transData.inverted().transform
        xlim = self.get_xlim(); xmid = sum(xlim)*0.5
        ylim = self.get_ylim(); ymid = sum(ylim)*0.5
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = xmid-x0; deltay = ymid-y0
        cx = cx or xmid; cy = cy or ymid
        xoff = (cx-xmid)*x
        self.set_xlim(xmid-deltax+xoff, xmid+deltax+xoff)
        yoff = (cy-ymid)*y
        self.set_ylim(ymid-deltay+yoff, ymid+deltay+yoff)

    def zoom(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view.
        """
        # get the function to convert display coordinates to data
        # coordinates
        transform = self.transData.inverted().transform
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = x0 - xlim[0]; deltay = y0 - ylim[0]
        self.set_xlim(xlim[0]+deltax, xlim[1]-deltax)
        self.set_ylim(ylim[0]+deltay, ylim[1]-deltay)

    def center_y(self, y):
        ymin, ymax = self.get_ylim()
        yoff = (ymax - ymin) * 0.5
        self.set_ylim(y-yoff, y+yoff)

    def center_x(self, x, offset=0.3):
        xmin, xmax = self.get_xlim()
        xspan = xmax - xmin
        xoff = xspan*0.5 + xspan*offset
        self.set_xlim(x-xoff, x+xoff)

    def scroll(self, x, y):
        x0, x1 = self.get_xlim()
        y0, y1 = self.get_ylim()
        xd = (x1-x0)*x
        yd = (y1-y0)*y
        self.set_xlim(x0+xd, x1+xd)
        self.set_ylim(y0+yd, y1+yd)

    def home(self):
        self.set_xlim(0, self.nchar)
        self.set_ylim(self.ntax, 0)

AlignmentPlot = subplot_class_factory(Alignment)

import os, matplotlib
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
from . import tree
import tempfile

## class TreeFigure:
##     def __init__(self):
##         pass

matplotlib.rcParams["xtick.direction"] = "out"

class TreeFigure:
    def __init__(self, root, relwidth=0.5, leafpad=1.5, name=None,
                 support=70.0, scaled=True, mark_named=True,
                 leaf_fontsize=10, branch_fontsize=10,
                 branch_width=1, branch_color="black",
                 highlight_support=True,
                 branchlabels=True, leaflabels=True, decorators=[],
                 xoff=0, yoff=0,
                 xlim=None, ylim=None,
                 height=None, width=None):
        self.root = root
        self.relwidth = relwidth
        self.leafpad = leafpad
        self.name = name
        self.support = support
        self.scaled = scaled
        self.mark_named = mark_named
        self.leaf_fontsize = leaf_fontsize
        self.branch_fontsize = branch_fontsize
        self.branch_width = branch_width
        self.branch_color = branch_color
        self.highlight_support = highlight_support
        self.branchlabels = branchlabels
        self.leaflabels = leaflabels
        self.decorators = decorators
        self.xoff = xoff
        self.yoff = yoff

        nleaves = len(root.leaves())
        self.dpi = 72.0
        h = height or (nleaves*self.leaf_fontsize*self.leafpad)/self.dpi
        self.height = h
        self.width = width or self.height*self.relwidth
        ## p = min(self.width, self.height)*0.1
        ## self.height += p
        ## self.width += p
        self.figure = Figure(figsize=(self.width, self.height), dpi=self.dpi)
        self.canvas = FigureCanvas(self.figure)
        self.axes = self.figure.add_axes(
            tree.TreePlot(self.figure, 1,1,1,
                          support=self.support,
                          scaled=self.scaled,
                          mark_named=self.mark_named,
                          leaf_fontsize=self.leaf_fontsize,
                          branch_fontsize=self.branch_fontsize,
                          branch_width=self.branch_width,
                          branch_color=self.branch_color,
                          highlight_support=self.highlight_support,
                          branchlabels=self.branchlabels,
                          leaflabels=self.leaflabels,
                          interactive=False,
                          decorators=self.decorators,
                          xoff=self.xoff, yoff=self.yoff,
                          name=self.name).plot_tree(self.root)
            )
        self.axes.spines["top"].set_visible(False)
        self.axes.spines["left"].set_visible(False)
        self.axes.spines["right"].set_visible(False)
        self.axes.spines["bottom"].set_smart_bounds(True)
        self.axes.xaxis.set_ticks_position("bottom")

        for v in self.axes.node2label.values():
            v.set_visible(True)

        ## for k, v in self.decorators:
        ##     func, args, kwargs = v
        ##     func(self.axes, *args, **kwargs)

        self.canvas.draw()
        ## self.axes.home()
        ## adjust_limits(self.axes)
        self.axes.set_position([0.05,0.05,0.95,0.95])

    @property
    def detail(self):
        return self.axes

    def savefig(self, fname):
        root, ext = os.path.splitext(fname)
        buf = tempfile.TemporaryFile()
        for i in range(3):
            self.figure.savefig(buf, format=ext[1:].lower())
            self.home()
            buf.seek(0)
        buf.close()
        self.figure.savefig(fname)

    def set_relative_width(self, relwidth):
        w, h = self.figure.get_size_inches()
        self.figure.set_figwidth(h*relwidth)

    def autoheight(self):
        "adjust figure height to show all leaf labels"
        nleaves = len(self.root.leaves())
        h = (nleaves*self.leaf_fontsize*self.leafpad)/self.dpi
        self.height = h
        self.figure.set_size_inches(self.width, self.height)
        self.axes.set_ylim(-2, nleaves+2)

    def home(self):
        self.axes.home()

"""
Convenience functions for drawing shapes on TreePlots.
"""
try:
    import Image
except ImportError:
    from PIL import Image
from numpy import pi
from matplotlib.collections import RegularPolyCollection, CircleCollection
from matplotlib.transforms import offset_copy
from matplotlib.patches import Rectangle, Wedge, Circle, PathPatch
from matplotlib.offsetbox import DrawingArea
from matplotlib.axes import Axes
from matplotlib.path import Path

try:
    from matplotlib.offsetbox import OffsetImage, AnnotationBbox
except ImportError:
    pass
from ..tree import Node
from . import colors as _colors

def _xy(plot, p):
    if isinstance(p, Node):
        c = plot.n2c[p]
        p = (c.x, c.y)
    elif isinstance(p, (list, tuple)):
        p = [ _xy(plot, x) for x in p ]
    else:
        pass
    return p

def image(plot, p, imgfile,
          maxdim=100, border=0,
          xoff=4, yoff=4,
          halign=0.0, valign=0.5,
          xycoords='data',
          boxcoords=('offset points')):
    """
    Add images to plot

    Args:
        plot (Tree): A Tree plot instance
        p (Node): A node object
        imgfile (str): A path to an image
        maxdim (float): Maximum dimension of image. Optional,
          defaults to 100.
        border: RR: What does border do? -CZ
        xoff, yoff (float): X and Y offset. Optional, defaults to 4
        halign, valign (float): Horizontal and vertical alignment within
          box. Optional, defaults to 0.0 and 0.5, respectively.

    """
    if xycoords == "label":
        xycoords = plot.node2label[p]
        x, y = (1, 0.5)
    else:
        x, y = _xy(plot, p)
    img = Image.open(imgfile)
    if max(img.size) > maxdim:
        img.thumbnail((maxdim, maxdim))
    imgbox = OffsetImage(img)
    abox = AnnotationBbox(imgbox, (x, y),
                          xybox=(xoff, yoff),
                          xycoords=xycoords,
                          box_alignment=(halign,valign),
                          pad=0.0,
                          boxcoords=boxcoords)
    plot.add_artist(abox)
    plot.figure.canvas.draw_idle()

def images(plot, p, imgfiles,
          maxdim=100, border=0,
          xoff=4, yoff=4,
          halign=0.0, valign=0.5,
          xycoords='data', boxcoords=('offset points')):
    """
    Add many images to plot at once

    Args:
        Plot (Tree): A Tree plot instance
        p (list): A list of node objects
        imgfile (list): A list of strs containing paths to image files.
          Must be the same length as p.
        maxdim (float): Maximum dimension of image. Optional,
          defaults to 100.
        border: RR: What does border do? -CZ
        xoff, yoff (float): X and Y offset. Optional, defaults to 4
        halign, valign (float): Horizontal and vertical alignment within
          box. Optional, defaults to 0.0 and 0.5, respectively.

    """
    for x, f in zip(p, imgfiles):
        image(plot, x, f, maxdim, border, xoff, yoff, halign, valign,
              xycoords, boxcoords)

def pie(plot, p, values, colors=None, size=16, norm=True,
        xoff=0, yoff=0,
        halign=0.5, valign=0.5,
        xycoords='data', boxcoords=('offset points')):
    """
    Draw a pie chart

    Args:
    plot (Tree): A Tree plot instance
    p (Node): A Node object
    values (list): A list of floats.
    colors (list): A list of strings to pull colors from. Optional.
    size (float): Diameter of the pie chart
    norm (bool): Whether or not to normalize the values so they
      add up to 360
    xoff, yoff (float): X and Y offset. Optional, defaults to 0
    halign, valign (float): Horizontal and vertical alignment within
      box. Optional, defaults to 0.5

    """
    x, y = _xy(plot, p)
    da = DrawingArea(size, size); r = size*0.5; center = (r,r)
    x0 = 0
    S = 360.0
    if norm: S = 360.0/sum(values)
    if not colors:
        c = _colors.tango()
        colors = [ c.next() for v in values ]
    for i, v in enumerate(values):
        theta = v*S
        if v: da.add_artist(Wedge(center, r, x0, x0+theta,
                                  fc=colors[i], ec='none'))
        x0 += theta
    box = AnnotationBbox(da, (x,y), pad=0, frameon=False,
                         xybox=(xoff, yoff),
                         xycoords=xycoords,
                         box_alignment=(halign,valign),
                         boxcoords=boxcoords)
    plot.add_artist(box)
    plot.figure.canvas.draw_idle()
    return box

def hbar(plot, p, values, colors=None, height=16,
         xoff=0, yoff=0,
         halign=1, valign=0.5,
         xycoords='data', boxcoords=('offset points')):
    x, y = _xy(plot, p)
    h = height; w = sum(values) * height#; yoff=h*0.5
    da = DrawingArea(w, h)
    x0 = -sum(values)
    if not colors:
        c = _colors.tango()
        colors = [ c.next() for v in values ]
    for i, v in enumerate(values):
        if v: da.add_artist(Rectangle((x0,0), v*h, h, fc=colors[i], ec='none'))
        x0 += v*h
    box = AnnotationBbox(da, (x,y), pad=0, frameon=False,
                         xybox=(xoff, yoff),
                         xycoords=xycoords,
                         box_alignment=(halign,valign),
                         boxcoords=boxcoords)
    plot.add_artist(box)
    plot.figure.canvas.draw_idle()

def hbars(plot, p, values, colors=None, height=16,
          xoff=0, yoff=0,
          halign=1, valign=0.5,
          xycoords='data', boxcoords=('offset points')):
    for x, v in zip(p, values):
        hbar(plot, x, v, colors, height, xoff, yoff, halign, valign,
             xycoords, boxcoords)

def squares(plot, p, colors='r', size=15, xoff=0, yoff=0, alpha=1.0,
            zorder=1000):
    """
    Draw a square at given node

    Args:
        plot (Tree): A Tree plot instance
        p: A node or list of nodes
        colors: Str or list of strs. Colors of squares to be drawn.
          Optional, defaults to 'r' (red)
        size (float): Size of the squares. Optional, defaults to 15
        xoff, yoff (float): Offset for x and y dimensions. Optional,
          defaults to 0.
        alpha (float): between 0 and 1. Alpha transparency of squares.
          Optional, defaults to 1 (fully opaque)
        zorder (int): The drawing order. Higher numbers appear on top
          of lower numbers. Optional, defaults to 1000.

    """
    points = _xy(plot, p)
    trans = offset_copy(
        plot.transData, fig=plot.figure, x=xoff, y=yoff, units='points')

    col = RegularPolyCollection(
        numsides=4, rotation=pi*0.25, sizes=(size*size,),
        offsets=points, facecolors=colors, transOffset=trans,
        edgecolors='none', alpha=alpha, zorder=zorder
        )

    plot.add_collection(col)
    plot.figure.canvas.draw_idle()

def tipsquares(plot, p, colors='r', size=15, pad=2, edgepad=10):
    """
    RR: Bug with this function. If you attempt to call it with a list as an
    argument for p, it will not only not work (expected) but it will also
    make it so that you can't interact with the tree figure (gives errors when
    you try to add symbols, select nodes, etc.) -CZ

    Add square after tip label, anchored to the side of the plot

    Args:
        plot (Tree): A Tree plot instance.
        p (Node): A Node object (Should be a leaf node).
        colors (str): olor of drawn square. Optional, defaults to 'r' (red)
        size (float): Size of square. Optional, defaults to 15
        pad: RR: I am unsure what this does. Does not seem to have visible
          effect when I change it. -CZ
        edgepad (float): Padding from square to edge of plot. Optional,
          defaults to 10.

    """
    x, y = _xy(plot, p) # p is a single node or point in data coordinates
    n = len(colors)
    da = DrawingArea(size*n+pad*(n-1), size, 0, 0)
    sx = 0
    for c in colors:
        sq = Rectangle((sx,0), size, size, color=c)
        da.add_artist(sq)
        sx += size+pad
    box = AnnotationBbox(da, (x, y), xybox=(-edgepad,y),
                         frameon=False,
                         pad=0.0,
                         xycoords='data',
                         box_alignment=(1, 0.5),
                         boxcoords=('axes points','data'))
    plot.add_artist(box)
    plot.figure.canvas.draw_idle()

def circles(plot, p, colors='g', size=15, xoff=0, yoff=0):
    """
    Draw circles on plot

    Args:
        plot (Tree): A Tree plot instance
        p: A node object or list of Node objects
        colors: Str or list of strs. Colors of the circles. Optional,
          defaults to 'g' (green)
        size (float): Size of the circles. Optional, defaults to 15
        xoff, yoff (float): X and Y offset. Optional, defaults to 0.

    """
    points = _xy(plot, p)
    trans = offset_copy(
        plot.transData, fig=plot.figure, x=xoff, y=yoff, units='points'
        )

    col = CircleCollection(
        sizes=(pi*size*size*0.25,),
        offsets=points, facecolors=colors, transOffset=trans,
        edgecolors='none'
        )

    plot.add_collection(col)
    plot.figure.canvas.draw_idle()
    return col

def legend(plot, colors, labels, shape='rectangle', loc='upper left', **kwargs):
    """
    RR: the MPL legend function has changed since this function has been
    written. This function currently does not work. -CZ
    """
    if shape == 'circle':
        shapes = [ Circle((0.5,0.5), radius=1, fc=c) for c in colors ]
        #shapes = [ CircleCollection([10],facecolors=[c]) for c in colors ]
    else:
        shapes = [ Rectangle((0,0),1,1,fc=c,ec='none') for c in colors ]

    return Axes.legend(plot, shapes, labels, loc=loc, **kwargs)

def leafspace_triangles(plot, color='black', rca=0.5):
    """
    RR: Using this function on the primates tree (straight from the newick file)
    gives error: 'Node' object has no attribute 'leafspace'. How do you give
    nodes the leafspace attribute? -CZ
    rca = relative crown age
    """
    leaves = plot.root.leaves()
    leafspace = [ float(x.leafspace) for x in leaves ]
    #leafspace = array(raw_leafspace)/(sum(raw_leafspace)/float(len(leaves)))
    pv = []
    for i, n in enumerate(leaves):
        if leafspace[i] > 0:
            p = plot.n2c[n]
            pp = plot.n2c[n.parent]
            spc = leafspace[i]
            yoff = spc/2.0
            x0 = pp.x + (p.x - pp.x)*rca
            verts = [(x0, p.y),
                     (p.x, p.y-yoff),
                     (p.x, p.y+yoff),
                     (x0, p.y)]
            codes = [Path.MOVETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY]
            path = Path(verts, codes)
            patch = PathPatch(path, fc=color, lw=0)
            pv.append(plot.add_patch(patch))
    return pv

def text(plot, x, y, s, color='black', xoff=0, yoff=0, valign='center',
         halign='left', fontsize=10):
    """
    Add text to the plot.

    Args:
        plot (Tree): A Tree plot instance
        x, y (float): x and y coordinates to place the text
        s (str): The text to write
        color (str): The color of the text. Optional, defaults to "black"
        xoff, yoff (float): x and y offset
        valign (str): Vertical alignment. Can be: 'center', 'top',
          'bottom', or 'baseline'. Defaults to 'center'.
        halign (str): Horizontal alignment. Can be: 'center', 'right',
          or 'left'. Defaults to 'left'
        fontsize (float): Font size. Optional, defaults to 10

    """
    txt = plot.annotate(
        s, xy=(x, y),
        xytext=(xoff, yoff),
        textcoords="offset points",
        verticalalignment=valign,
        horizontalalignment=halign,
        fontsize=fontsize,
        clip_on=True,
        picker=True
    )
    txt.set_visible(True)
    return txt

"""
interactive viewers for trees, etc. using matplotlib
"""
import sys, math, types, os, operator
from itertools import chain
from .. import tree
from ..layout import cartesian
from ..storage import Storage
from .. import pyperclip as clipboard
# from ..nodecache import NodeCache
import matplotlib.pyplot as pyplot
from matplotlib.axes import Axes, subplot_class_factory
from matplotlib.figure import SubplotParams
from matplotlib.patches import PathPatch, Rectangle, Arc
from matplotlib.path import Path
from matplotlib.widgets import RectangleSelector
from matplotlib import cm as mpl_colormap
from matplotlib import colors as mpl_colors
from matplotlib.collections import LineCollection
try:
    from matplotlib.offsetbox import OffsetImage, AnnotationBbox
except ImportError:
    pass
from matplotlib.ticker import NullLocator
from mpl_toolkits.axes_grid.anchored_artists import AnchoredText
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from . import symbols, colors
from . import hardcopy as HC
try:
    import Image
except ImportError:
    from PIL import Image

# matplotlib.rcParams['path.simplify'] = False

_tango = colors.tango()
class TreeFigure(object):
    """
    Window for showing a single tree, optionally with split overview
    and detail panes.

    The navigation toolbar at the bottom is provided by matplotlib
    (http://matplotlib.sf.net/users/navigation_toolbar.html).  Its
    pan/zoom button and zoom-rectangle button provide different modes
    of mouse interaction with the figure.  When neither of these
    buttons are checked, the default mouse bindings are as follows:

    * button 1 drag: select nodes - retrieve by calling fig.selected_nodes
    * button 3 drag: pan view
    * scroll up/down: zoom in/out
    * scroll up/down with Control key: zoom y-axis
    * scroll up/down with Shift key: zoom x-axis
    * scroll up/down with 'd' key: pan view up/down
    * scroll up/down with 'e' key: pan view left/right
    * click on overview will center the detail pane on that region

    Default keybindings:

    * t: zoom out to full extent
    * +/-: zoom in/out

    Useful attributes and methods (assume an instance named *fig*):

    * fig.root - the root node (see [Node methods])
    * fig.highlight(s) - highlight and trace nodes with substring *s*
    * fig.zoom_clade(anc) - zoom to view node *anc* and all its descendants
    * fig.toggle_overview() - toggle visibility of the overview pane
    * fig.toggle_branchlabels() - ditto for branch labels
    * fig.toggle_leaflabels() - ditto for leaf labels
    * fig.decorate(func) - decorate the tree with a function (see
      :ref:`decorating TreeFigures <decorating-trees>`)
    """
    def __init__(self, data, name=None, scaled=True, div=0.25,
                 branchlabels=True, leaflabels=True, mark_named=True,
                 highlight_support=True, xoff=0, yoff=0,
                 overview=True, radial=False):
        self.overview = None
        self.overview_width = div
        self.dataplot = None
        self.dataplot_width = 0.25
        self.name = name
        self.scaled = scaled
        self.branchlabels = branchlabels
        self.leaflabels = leaflabels
        self.mark_named = mark_named
        self.xoff = xoff
        self.yoff = yoff
        self.radial = radial
        if radial:
            self.leaflabels = False
        self.highlighted = set()
        self.highlight_support = highlight_support
        if isinstance(data, tree.Node):
            root = data
        else:
            root = tree.read(data)
        self.root = root
        if not self.root:
            raise IOError("cannot coerce data into tree.Node")
        self.name = self.name or root.treename
        pars = SubplotParams(
            left=0, right=1, bottom=0.05, top=1, wspace=0.01
            )
        fig = pyplot.figure(subplotpars=pars, facecolor="white")
        connect_events(fig.canvas)
        self.figure = fig
        self.initialize_subplots(overview)
        self.home()

    def initialize_subplots(self, overview=True):
        if not self.radial:
            tp = TreePlot(self.figure, 1, 2, 2, app=self, name=self.name,
                          scaled=self.scaled, branchlabels=self.branchlabels,
                          highlight_support=self.highlight_support,
                          leaflabels=self.leaflabels,
                          mark_named=self.mark_named)
            detail = self.figure.add_subplot(tp)
            detail.set_root(self.root)
            detail.plot_tree()
            self.detail = detail
            tp = OverviewTreePlot(
                self.figure, 121, app=self, scaled=self.scaled,
                branchlabels=False, leaflabels=False,
                mark_named=self.mark_named,
                highlight_support=self.highlight_support,
                target=self.detail
                )
            ov = self.figure.add_subplot(tp)
            ov.set_root(self.root)
            ov.plot_tree()
            self.overview = ov
            if not overview:
                self.toggle_overview(False)
            self.set_positions()

            if self.detail.nleaves < 50:
                self.toggle_overview(False)
        else:
            tp = RadialTreePlot(
                self.figure, 111, app=self, name=self.name,
                scaled=self.scaled, branchlabels=self.branchlabels,
                highlight_support=self.highlight_support,
                leaflabels=self.leaflabels, mark_named=self.mark_named
                )
            ax2 = self.figure.add_subplot(tp)
            ax2.set_root(self.root)
            ax2.plot_tree()
            self.detail = ax2

    def __get_selected_nodes(self):
        return list(self.detail.selected_nodes)

    def __set_selected_nodes(self, nodes):
        self.detail.select_nodes(nodes)

    def __del_selected_nodes(self):
        self.detail.select_nodes(None)

    selected = property(__get_selected_nodes,
                        __set_selected_nodes,
                        __del_selected_nodes)

    ## def selected_nodes(self):
    ##     return self.detail.selected_nodes

    @property
    def axes(self):
        return self.detail

    def add(self, data, name=None, support=70,
            branchlabels=False, leaflabels=True, mark_named=True):
        """
        Add a new tree in a new window

        Args:
            data: A node object or tree file.
            name (str): Name of the plot. Defaults to None
            branchlabels (bool): Whether or not to draw branch labels.
              Defaults to False
            leaflabels (bool): Whether or not to draw leaf labels.
              Defaults to True
        """
        newfig = MultiTreeFigure()
        ## newfig.add(self.root, name=self.name, support=self.support,
        ##            branchlabels=self.branchlabels)
        newfig.add(data, name=name, support=support,
                   branchlabels=branchlabels,
                   leaflabels=leaflabels,
                   mark_named=mark_named)
        return newfig

    def toggle_leaflabels(self):
        """
        Toggle leaf labels and redraw tree
        """
        self.leaflabels = not self.leaflabels
        self.detail.leaflabels = self.leaflabels
        self.redraw()

    def toggle_branchlabels(self):
        """
        Toggle branch labels and redraw tree
        """
        self.branchlabels = not self.branchlabels
        self.detail.branchlabels = self.branchlabels
        self.redraw()

    def toggle_overview(self, val=None):
        """
        Toggle overview
        """
        if val is None:
            if self.overview.get_visible():
                self.overview.set_visible(False)
                self.overview_width = 0.001
            else:
                self.overview.set_visible(True)
                self.overview_width = 0.25
        elif val:
            self.overview.set_visible(True)
            self.overview_width = val
        else:
            self.overview.set_visible(False)
            self.overview_width = 0.001
        self.set_positions()

    def set_scaled(self, scaled):
        """
        RR: Using this method gives the error:
            redraw takes exactly 1 argument(2 given)-CZ
        Define whether or not the tree is scaled and redraw tree

        Args:
            scaled (bool): Whether or not the tree is scaled.
        """
        for p in self.overview, self.detail:
            p.redraw(p.set_scaled(scaled))
        self.set_positions()

    def on_nodes_selected(self, treeplot):
        pass

    def picked(self, e):
        try:
            if e.mouseevent.button==1:
                s = e.artist.get_text()
                clipboard.copy(s)
                print(s)
                sys.stdout.flush()
        except:
            pass

    def ladderize(self, rev=False):
        """
        Ladderize and redraw the tree
        """
        self.root.ladderize(rev)
        self.redraw()

    def show(self):
        """
        Plot the figure in a new window
        """
        self.figure.show()

    def set_positions(self):
        ov = self.overview
        p = self.detail
        dp = self.dataplot
        height = 1.0-p.xoffset()
        if ov:
            box = [0, p.xoffset(), self.overview_width, height]
            ov.set_position(box)
        w = 1.0
        if ov:
            w -= self.overview_width
        if dp:
            w -= self.dataplot_width
        p.set_position([self.overview_width, p.xoffset(), w, height])
        if dp:
            box = [1.0-self.dataplot_width, p.xoffset(),
                   self.dataplot_width, height]
            dp.set_position(box)
        self.figure.canvas.draw_idle()

    ## def div(self, v=0.3):
    ##     assert 0 <= v < 1
    ##     self.overview_width = v
    ##     self.set_positions()
    ##     self.figure.canvas.draw_idle()

    def add_dataplot(self):
        """
        Add new plot to the side of existing plot
        """
        np = 3 if self.overview else 2
        if self.dataplot:
            self.figure.delaxes(self.dataplot)
        self.dataplot = self.figure.add_subplot(1, np, np, sharey=self.detail)
        # left, bottom, width, height (proportions)
        dleft, dbottom, dwidth, dheight = self.detail.get_position().bounds
        # give the dataplot one-quarter the width of the detail axes
        w = dwidth * 0.25
        self.detail.set_position([dleft, dbottom, dwidth-w, dheight])
        self.dataplot.set_position([1-w, dbottom, w, dheight])
        self.dataplot.xaxis.set_visible(False)
        self.dataplot.yaxis.set_visible(False)
        for x in self.dataplot.spines.values():
            x.set_visible(False)
        self.figure.canvas.draw_idle()
        return self.dataplot

    def redraw(self):
        """
        Replot the figure and overview
        """
        self.detail.redraw()
        if self.overview: self.overview.redraw()
        self.highlight()
        self.set_positions()
        self.figure.canvas.draw_idle()

    def find(self, x):
        """
        Find nodes

        Args:
            x (str): String to search
        Returns:
            list: A list of node objects found with the Node findall() method
        """
        return self.root.findall(x)

    def hlines(self, nodes, width=5, color="red", xoff=0, yoff=0):
        """
        Highlight nodes

        Args:
            nodes (list): A list of node objects
            width (float): Width of highlighted lines. Defaults to 5
            color (str): Color of highlighted lines. Defaults to red
            xoff (float): Number of units to offset lines by. Defaults to 0
            yoff (float): Number of units to offset lines by. Defaults to 0
        """
        self.overview.hlines(nodes, width=width, color=color,
                             xoff=xoff, yoff=yoff)
        self.detail.hlines(nodes, width=width, color=color,
                           xoff=xoff, yoff=yoff)

    def highlight(self, x=None, width=5, color="red"):
        """
        Highlight nodes

        Args:
            x: Str or list of Strs or Node or list of Nodes
            width (float): Width of highlighted lines. Defaults to 5
            color (str): Color of highlighted lines. Defaults to red
        """
        if x:
            nodes = set()
            if type(x) in types.StringTypes:
                nodes = self.root.findall(x)
            elif isinstance(x, tree.Node):
                nodes = set(x)
            else:
                for n in x:
                    if type(n) in types.StringTypes:
                        found = self.root.findall(n)
                        if found:
                            nodes |= set(found)
                    elif isinstance(n, tree.Node):
                        nodes.add(n)

            self.highlighted = nodes
        else:
            self.highlighted = set()
        if self.overview:
            self.overview.highlight(self.highlighted, width=width, color=color)
        self.detail.highlight(self.highlighted, width=width, color=color)
        self.figure.canvas.draw_idle()

    def home(self):
        """
        Return plot to initial size and location.
        """
        if self.overview: self.overview.home()
        self.detail.home()

    def zoom_clade(self, x):
        """
        Zoom to fit a node *x* and all its descendants in the view.

        Args:
            x: Node or str that matches the label of a node
        """
        if not isinstance(x, tree.Node):
            x = self.root[x]
        self.detail.zoom_clade(x)

    def zoom(self, factor=0.1):
        """Zoom both axes by *factor* (relative display size)."""
        self.detail.zoom(factor, factor)
        self.figure.canvas.draw_idle()

    def zx(self, factor=0.1):
        """Zoom x axis by *factor*."""
        self.detail.zoom(factor, 0)
        self.figure.canvas.draw_idle()

    def zy(self, factor=0.1):
        """Zoom y axis by *factor*."""
        self.detail.zoom(0, factor)
        self.figure.canvas.draw_idle()

    def decorate(self, func, *args, **kwargs):
        """
        Decorate the tree.

        Args:
            func (function): A function that takes a TreePlot instance as the
              first parameter, and *args* and *kwargs* as additional
              parameters.  It adds boxes, circles, etc to the TreePlot.

        Notes:
            If *kwargs* contains the key-value pair ('store', *name*),
            then the function is stored as *name* and re-called every time
            the TreePlot is redrawn, i.e., the decoration is persistent.
            Use ``rmdec(name)`` to remove the decorator from the treeplot.
        """
        self.detail.decorate(func, *args, **kwargs)

    def rmdec(self, name):
        "Remove the decoration 'name'."
        self.detail.rmdec(name)
        ## if name in self.detail.decorators:
        ##     del self.detail.decorators[name]

    def cbar(self, node, width=6, color='blue', mrca = True):
        pass
        # self.axes.cbar(nodes = node, width = width, color = color, mrca = mrca)

    def unclutter(self, *args):
        self.detail.unclutter()

    def trace_branches(self, nodes, width=4, color="blue"):
        """
        RR: What is the difference between this and highlight? -CZ
        """
        for p in self.overview, self.detail:
            p.trace_branches(nodes, width, color)

    def plot_continuous(self, *args, **kwargs):
        self.detail.plot_continuous(*args, **kwargs)

    def hardcopy(self, fname=None, relwidth=None, leafpad=1.5):
        if not relwidth:
            bbox = self.detail.get_tightbbox(self.figure.canvas.get_renderer())
            relwidth = bbox.width/bbox.height
        f = self.detail.hardcopy(
            relwidth=relwidth,
            leafpad=leafpad
            )
        f.axes.home()
        #f.axes.set_xlim(*self.detail.get_xlim())
        #f.axes.set_ylim(*self.detail.get_ylim())
        if fname:
            f.savefig(fname)
        return f

    def select_nodes(self, nodes=None):
        """
        Select nodes on the plot

        Args:
            nodes: A node or list of ndoes
        Notes:
            If only one node is given, all of the node's ancestors are
            also selected. If a list of nodes is given (even if it has only
            one node), only the given node(s) are selected.
        """
        self.detail.select_nodes(nodes)

    def decorate(self, func, *args, **kwargs): # RR: is this repeated from above? -CZ
        self.detail.decorate(func, *args, **kwargs)

    ## def dataplot(self):
    ##     ax = self.figure.add_subplot(133, sharey=self.detail)
    ##     ax.yaxis.set_visible(False)
    ##     self.dataplot = ax
    ##     return ax

    def attach_alignment(self, aln, overview=True):
        "leaf labels expected to be sequence ids"
        from Bio.Align import MultipleSeqAlignment
        from Bio.Seq import Seq
        from Bio.SeqRecord import SeqRecord
        from Bio.Alphabet import IUPAC
        from alignment import AlignmentFigure, AlignmentPlot
        if not isinstance(aln, MultipleSeqAlignment):
            from .. import align
            aln = align.read(aln)
        d = dict([ (x.id,x) for x in aln ])
        emptyseq = Seq('-'*aln.get_alignment_length(),
                       alphabet=IUPAC.ambiguous_dna)
        aln = MultipleSeqAlignment(
            [ d.get(x.label) or SeqRecord(emptyseq, id=x.label)
              for x in self.root.leaves() ]
            )
        self.aln = aln
        p = AlignmentPlot(self.figure, 133, aln=aln, app=self,
                          sharey=self.detail, showy=False)
        self.alnplot = Storage()
        self.alnplot.detail = self.figure.add_subplot(p)
        detail = self.alnplot.detail
        detail.plot_aln()
        if overview:
            self.alnplot.overview = inset_axes(
                detail, width="30%", height="20%", loc=1
                )
            overview = self.alnplot.overview
            overview.xaxis.set_major_locator(NullLocator())
            overview.yaxis.set_major_locator(NullLocator())
            overview.imshow(
                detail.array, interpolation='nearest', aspect='auto',
                origin='lower'
                )
            rect = UpdatingRect(
                [0,0], 0, 0, facecolor='black', edgecolor='cyan', alpha=0.5
                )
            overview.zoomrect = rect
            rect.target = detail
            detail.callbacks.connect('xlim_changed', rect)
            detail.callbacks.connect('ylim_changed', rect)
            overview.add_patch(rect)
            rect(overview)
        self.toggle_overview(False)
        xoff = self.detail.xoffset()
        self.detail.set_position([0, xoff, 0.3, 1.0-xoff])
        p.set_position([0.3, xoff, 0.7, 1.0-xoff])

class MultiTreeFigure(object):
    """
    Window for showing multiple trees side-by-side.

    TODO: document this
    """
    def __init__(self, trees=None, name=None, support=70,
                 scaled=True, branchlabels=False, radial=False):
        """
        *trees* are assumed to be objects suitable for passing to
        ivy.tree.read()
        """
        self.root = []
        self.name = name
        self.name2plot = {}
        self.plot = []
        self.scaled = scaled
        self.branchlabels = branchlabels
        self.radial = radial
        self.highlighted = set()
        self.divs = []
        pars = SubplotParams(
            left=0, right=1, bottom=0.05, top=1, wspace=0.04
            )
        fig = pyplot.figure(subplotpars=pars)
        connect_events(fig.canvas)
        self.figure = fig

        for x in trees or []:
            self.add(x, support=support, scaled=scaled,
                     branchlabels=branchlabels)

    def on_nodes_selected(self, treeplot):
        pass

    def clear(self):
        self.root = []
        self.name2plot = {}
        self.highlighted = set()
        self.divs = []
        self.figure.clf()

    def picked(self, e):
        try:
            if e.mouseevent.button==1:
                print(e.artist.get_text())
                sys.stdout.flush()
        except:
            pass

    def getplot(self, x):
        p = None
        try:
            i = self.root.index(x)
            return self.plot[i]
        except ValueError:
            return self.name2plot.get(x)

    def add(self, data, name=None, support=70, scaled=True,
            branchlabels=False, leaflabels=True, mark_named=True):
        root = None
        if isinstance(data, tree.Node):
            root = data
        else:
            root = tree.read(data)
        if not root:
            raise IOError("cannot coerce data into tree.Node")

        name = name or root.treename
        self.root.append(root)

        fig = self.figure
        N = len(self.plot)+1
        for i, p in enumerate(self.plot):
            p.change_geometry(1, N, i+1)
        plt = TreePlot(fig, 1, N, N, app=self, name=name, support=support,
                       scaled=scaled, branchlabels=branchlabels,
                       leaflabels=leaflabels, mark_named=mark_named)
        p = fig.add_subplot(plt)
        p.set_root(root)
        p.plot_tree()
        p.index = N-1
        self.plot.append(p)
        if name:
            assert name not in self.name2plot
            self.name2plot[name] = p

        ## global IP
        ## if IP:
        ##     def f(shell, s):
        ##         self.highlight(s)
        ##         return sorted([ x.label for x in self.highlighted ])
        ##     IP.expose_magic("highlight", f)
        ##     def f(shell, s):
        ##         self.root.ladderize()
        ##         self.redraw()
        ##     IP.expose_magic("ladderize", f)
        ##     def f(shell, s):
        ##         self.show()
        ##     IP.expose_magic("show", f)
        ##     def f(shell, s):
        ##         self.redraw()
        ##     IP.expose_magic("redraw", f)
        return p

    def show(self):
        self.figure.show()

    def redraw(self):
        for p in self.plot:
            p.redraw()
        self.figure.canvas.draw_idle()

    def ladderize(self, reverse=False):
        for n in self.root:
            n.ladderize(reverse)
        self.redraw()

    def highlight(self, s=None, add=False, width=5, color="red"):
        """
        Highlight nodes

        Args:
            s: Str or list of Strs or Node or list of Nodes
            add (bool): Whether to add to existing highlighted nodes or
              overwrite them.
            width (float): Width of highlighted lines. Defaults to 5
            color (str): Color of highlighted lines. Defaults to red
        """
        if not s:
            self.highlighted = set()
        if not add:
            self.highlighted = set()

        nodesets = [ p.root.findall(s) for p in self.plot ]

        for nodes, p in zip(nodesets, self.plot):
            if nodes:
                p.highlight(nodes, width=width, color=color)
            else:
                p.highlight()

        self.highlighted = nodesets
        self.figure.canvas.draw_idle()

        ##     for root in self.root:
        ##         for node in root.iternodes():
        ##             if node.label and (s in node.label):
        ##                 self.highlighted.add(node)
        ## self.highlight()

    def home(self):
        for p in self.plot: p.home()

def connect_events(canvas):
    mpl_connect = canvas.mpl_connect
    mpl_connect("button_press_event", onclick)
    mpl_connect("button_release_event", onbuttonrelease)
    mpl_connect("scroll_event", onscroll)
    mpl_connect("pick_event", onpick)
    mpl_connect("motion_notify_event", ondrag)
    mpl_connect("key_press_event", onkeypress)
    mpl_connect("axes_enter_event", axes_enter)
    mpl_connect("axes_leave_event", axes_leave)

class UpdatingRect(Rectangle):
    def __call__(self, p):
        self.set_bounds(*p.viewLim.bounds)
        p.figure.canvas.draw_idle()

class Tree(Axes):
    """
    matplotlib.axes.Axes subclass for rendering trees.
    """
    def __init__(self, fig, rect, *args, **kwargs):
        self.root = None
        self.app = kwargs.pop("app", None)
        self.support = kwargs.pop("support", 70.0)
        self.scaled = kwargs.pop("scaled", True)
        self.leaflabels = kwargs.pop("leaflabels", True)
        self.branchlabels = kwargs.pop("branchlabels", True)
        self._mark_named = kwargs.pop("mark_named", True)
        self.name = None
        self.leaf_fontsize = kwargs.pop("leaf_fontsize", 10)
        self.branch_fontsize = kwargs.pop("branch_fontsize", 10)
        self.branch_width = kwargs.pop("branch_width", 1)
        self.branch_color = kwargs.pop("branch_color", "black")
        self.interactive = kwargs.pop("interactive", True)
        self.decorators = kwargs.pop("decorators", [])
        ## if self.decorators:
        ##     print >> sys.stderr, "got %s decorators" % len(self.decorators)
        self.xoff = kwargs.pop("xoff", 0)
        self.yoff = kwargs.pop("yoff", 0)
        self.highlight_support = kwargs.pop("highlight_support", True)
        self.smooth_xpos = kwargs.pop("smooth_xpos", 0)
        Axes.__init__(self, fig, rect, *args, **kwargs)
        self.nleaves = 0
        self.highlighted = None
        self.highlightpatch = None
        self.pan_start = None
        if not self.decorators:
            self.decorators = [
                ("__selected_nodes__", (Tree.highlight_selected_nodes, [], {}))
                ]
        self.name2dec = dict([ (x[0], i) for i, x in
                               enumerate(self.decorators) ])
        self._active = False

        if self.interactive:
            self.callbacks.connect("ylim_changed", self.draw_labels)
        self.selector = RectangleSelector(self, self.rectselect,
                                          useblit=True)
        def f(e):
            if e.button != 1: return True
            else: return RectangleSelector.ignore(self.selector, e)
        self.selector.ignore = f
        self.xoffset_value = 0.05
        self.selected_nodes = set()
        self.leaf_offset = 4
        self.leaf_valign = "center"
        self.leaf_halign = "left"
        self.branch_offset = -5
        self.branch_valign = "center"
        self.branch_halign = "right"

        self.spines["top"].set_visible(False)
        self.spines["left"].set_visible(False)
        self.spines["right"].set_visible(False)
        self.xaxis.set_ticks_position("bottom")

    def p2y(self):
        "Convert a single display point to y-units"
        transform = self.transData.inverted().transform
        return transform([0,1])[1] - transform([0,0])[1]

    def p2x(self):
        "Convert a single display point to y-units"
        transform = self.transData.inverted().transform
        return transform([0,0])[1] - transform([1,0])[1]

    def decorate(self, func, *args, **kwargs):
        """
        Decorate the tree with function *func*.  If *kwargs* contains
        the key-value pair ('store', *name*), the decorator function
        is stored in self.decorators and called upon every redraw.
        """
        name = kwargs.pop("store", None)
        if name:
            if name in self.name2dec:
                i = self.name2dec[name]
                self.decorators[i] = (name, (func, args, kwargs))
            else:
                self.decorators.append((name, (func, args, kwargs)))
                self.name2dec = dict([ (x[0], i) for i, x in
                                       enumerate(self.decorators) ])

        func(self, *args, **kwargs)

    def rmdec(self, name):
        if name in self.name2dec:
            i = self.name2dec[name]
            del self.decorators[i]
            self.name2dec = dict([ (x[0], i) for i, x in
                                   enumerate(self.decorators) ])

    def flip(self):
        """
        Reverse the direction of the x-axis.
        """
        self.leaf_offset *= -1
        self.branch_offset *= -1
        ha = self.leaf_halign
        self.leaf_halign = "right" if ha == "left" else "left"
        ha = self.branch_halign
        self.branch_halign = "right" if ha == "left" else "left"
        self.invert_xaxis()
        self.redraw()

    def xoffset(self):
        """Space below x axis to show tick labels."""
        if self.scaled:
            return self.xoffset_value
        else:
            return 0

    def save_newick(self, filename):
        """
        Save tree as a newick file.

        Args:
            filename (str): Path to file.

        """
        if os.path.exists(filename):
            s = raw_input("File %s exists, enter 'y' to overwrite ").strip()
            if (s and s.lower() != 'y') or (not s):
                return
        import newick
        f = file(filename, "w")
        f.write(newick.string(self.root))
        f.close()

    def set_scaled(self, scaled):
        flag = self.scaled != scaled
        self.scaled = scaled
        return flag

    def cbar(self, nodes, color=None, label=None, x=None, width=8, xoff=10,
             showlabel=True, mrca=True):
        """
        Draw a 'clade' bar (i.e., along the y-axis) indicating a
        clade.  *nodes* are assumed to be one or more nodes in the
        tree.  If just one, it should be the internal node
        representing the clade of interest; otherwise, the clade of
        interest is the most recent common ancestor of the specified
        nodes.  *label* is an optional string to be drawn next to the
        bar, *offset* by the specified number of display units.  If
        *label* is ``None`` then the clade's label is used instead.

        Args:
            nodes: Node or list of nodes
            color (str): Color of the bar. Optional, defaults to None.
            label (str): Optional label for bar. If None, the clade's
              label is used instead. Defaults to None.
            width (float): Width of bar
            xoff (float): Offset from label to bar
            showlabel (bool): Whether or not to draw the label
            mrca: RR: Not quite sure what this does -CZ

        """
        xlim = self.get_xlim(); ylim = self.get_ylim()
        if color is None: color = _tango.next()
        transform = self.transData.inverted().transform

        if mrca:
            if isinstance(nodes, tree.Node):
                spec = nodes
            elif type(nodes) in types.StringTypes:
                spec = self.root.get(nodes)
            else:
                spec = self.root.mrca(nodes)

            assert spec in self.root
            label = label or spec.label
            leaves = spec.leaves()

        else:
            leaves = nodes

        n2c = self.n2c

        y = sorted([ n2c[n].y for n in leaves ])
        ymin = y[0]; ymax = y[-1]; y = (ymax+ymin)*0.5

        if x is None:
            x = max([ n2c[n].x for n in leaves ])
            _x = 0
            for lf in leaves:
                txt = self.node2label.get(lf)
                if txt and txt.get_visible():
                    _x = max(_x, transform(txt.get_window_extent())[1,0])
            if _x > x: x = _x

        v = sorted(list(transform(((0,0),(xoff,0)))[:,0]))
        xoff = v[1]-v[0]
        x += xoff

        Axes.plot(self, [x,x], [ymin, ymax], '-', linewidth=width, color=color)

        if showlabel and label:
            xo = self.leaf_offset
            if xo > 0:
                xo += width*0.5
            else:
                xo -= width*0.5
            txt = self.annotate(
                label,
                xy=(x, y),
                xytext=(xo, 0),
                textcoords="offset points",
                verticalalignment=self.leaf_valign,
                horizontalalignment=self.leaf_halign,
                fontsize=self.leaf_fontsize,
                clip_on=True,
                picker=False
                )

        self.set_xlim(xlim); self.set_ylim(ylim)

    def anctrace(self, anc, descendants=None, width=4, color="blue"):
        """
        RR: This function gives me a 'list index out of range' error
        when I try to use it -CZ
        """
        if not descendants:
            descendants = anc.leaves()
        else:
            for d in descendants:
                assert d in anc

        nodes = []
        for d in descendants:
            v = d.rootpath(anc)
            if v:
                nodes.extend(v)
        nodes = set(nodes)
        nodes.remove(anc)
        self.trace_branches(nodes, width, color)

    def trace_branches(self, nodes, width=4, color="blue"):
        n2c = self.n2c
        M = Path.MOVETO; L = Path.LINETO
        verts = []
        codes = []
        for c, pc in [ (n2c[x], n2c[x.parent]) for x in nodes
                       if (x in n2c) and x.parent ]:
            x = c.x; y = c.y
            px = pc.x; py = pc.y
            verts.append((x, y)); codes.append(M)
            verts.append((px, y)); codes.append(L)
            verts.append((px, py)); codes.append(L)
        px, py = verts[-1]
        verts.append((px, py)); codes.append(M)

        p = PathPatch(Path(verts, codes), fill=False,
                      linewidth=width, edgecolor=color)
        self.add_patch(p)
        self.figure.canvas.draw_idle()
        return p

    def highlight_selected_nodes(self, color="green"):
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        get = self.n2c.get
        coords = list(filter(None, [ get(n) for n in self.selected_nodes ]))
        x = [ c.x for c in coords ]
        y = [ c.y for c in coords ]
        if x and y:
            self.__selected_highlight_patch = self.scatter(x, y, s=60, c=color,
                                                           zorder=100)
        self.set_xlim(xlim)
        self.set_ylim(ylim)
        self.figure.canvas.draw_idle()

    def select_nodes(self, nodes=None, add=False):
        try:
            self.__selected_highlight_patch.remove()
            self.figure.canvas.draw_idle()
        except:
            pass
        if add:
            if nodes:
                self.selected_nodes = self.selected_nodes | nodes
            if hasattr(self, "app") and self.app:
                self.app.on_nodes_selected(self)
            self.highlight_selected_nodes()
        else:
            if nodes:
                self.selected_nodes = nodes
                if hasattr(self, "app") and self.app:
                    self.app.on_nodes_selected(self)
                self.highlight_selected_nodes()
            else:
                self.selected_nodes = set()

    def rectselect(self, e0, e1):
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        s = set()
        x0, x1 = sorted((e0.xdata, e1.xdata))
        y0, y1 = sorted((e0.ydata, e1.ydata))
        add = e0.key == 'shift'
        for n, c in self.n2c.items():
            if (x0 < c.x < x1) and (y0 < c.y < y1):
                s.add(n)
        self.select_nodes(nodes = s, add = add)
        self.set_xlim(xlim)
        self.set_ylim(ylim)
        ## if s:
        ##     print "Selected:"
        ##     for n in s:
        ##         print " ", n

    def picked(self, e):
        if hasattr(self, "app") and self.app:
            self.app.picked(e)

    def window2data(self, expandx=1.0, expandy=1.0):
        """
        return the data coordinates ((x0, y0),(x1, y1)) of the plot
        window, expanded by relative units of window size
        """
        bb = self.get_window_extent()
        bbx = bb.expanded(expandx, expandy)
        return self.transData.inverted().transform(bbx.get_points())

    def get_visible_nodes(self, labeled_only=False):
        ## transform = self.transData.inverted().transform
        ## bb = self.get_window_extent()
        ## bbx = bb.expanded(1.1,1.1)
        ## ((x0, y0),(x1, y1)) = transform(bbx.get_points())
        ((x0, y0),(x1, y1)) = self.window2data(1.1, 1.1)
        #print "visible_nodes points", x0, x1, y0, y1

        if labeled_only:
            def f(v): return (y0 < v[0] < y1) and (v[2] in self.node2label)
        else:
            def f(v): return (y0 < v[0] < y1)
        for y, x, n in filter(f, self.coords):
            yield (n, x, y)

    def zoom_cxy(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view, with a fixed data point (cx, cy)
        """
        transform = self.transData.inverted().transform
        xlim = self.get_xlim(); xmid = sum(xlim)*0.5
        ylim = self.get_ylim(); ymid = sum(ylim)*0.5
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = xmid-x0; deltay = ymid-y0
        cx = cx or xmid; cy = cy or ymid
        xoff = (cx-xmid)*x
        self.set_xlim(xmid-deltax+xoff, xmid+deltax+xoff)
        yoff = (cy-ymid)*y
        self.set_ylim(ymid-deltay+yoff, ymid+deltay+yoff)
        self.adjust_xspine()

    def zoom(self, x=0.1, y=0.1, cx=None, cy=None):
        """
        Zoom the x and y axes in by the specified proportion of the
        current view.
        """
        # get the function to convert display coordinates to data
        # coordinates
        transform = self.transData.inverted().transform
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        bb = self.get_window_extent()
        bbx = bb.expanded(1.0-x,1.0-y)
        points = transform(bbx.get_points())
        x0, x1 = points[:,0]; y0, y1 = points[:,1]
        deltax = x0 - xlim[0]; deltay = y0 - ylim[0]
        self.set_xlim(xlim[0]+deltax, xlim[1]-deltax)
        self.set_ylim(ylim[0]+deltay, ylim[1]-deltay)
        self.adjust_xspine()

    def center_y(self, y):
        """
        Center the y-axis of the canvas on the given y value
        """
        ymin, ymax = self.get_ylim()
        yoff = (ymax - ymin) * 0.5
        self.set_ylim(y-yoff, y+yoff)
        self.adjust_xspine()

    def center_x(self, x, offset=0.3):
        """
        Center the x-axis of the canvas on the given x value
        """
        xmin, xmax = self.get_xlim()
        xspan = xmax - xmin
        xoff = xspan*0.5 + xspan*offset
        self.set_xlim(x-xoff, x+xoff)
        self.adjust_xspine()

    def center_node(self, node):
        """
        Center the canvas on the given node
        """
        c = self.n2c[node]
        y = c.y
        self.center_y(y)
        x = c.x
        self.center_x(x, 0.2)

    def do_highlight_support(self):
        """
        TODO: reconfigure this, insert into self.decorators
        """
        if self.support:
            lim = float(self.support)

        M = Path.MOVETO; L = Path.LINETO

        verts = []; codes = []
        segments = []
        def f(n):
            if n.isleaf or not n.parent: return False
            try: return float(n.label) >= lim
            except:
                try: return float(n.support) >= lim
                except: pass
            return False

        for node, coords in [ x for x in self.n2c.items() if f(x[0]) ]:
            x = coords.x; y = coords.y
            p = node.parent
            pcoords = self.n2c[p]
            px = pcoords.x; py = y
            if self.app and self.app.radial:
                pc = self.n2c[node.parent]; theta2 = pc.angle
                px = math.cos(math.radians(coords.angle))*pc.depth
                py = math.sin(math.radians(coords.angle))*pc.depth

            ## segments.append([(x, y),(px, y)])
            verts.append((x,y)); codes.append(M)
            verts.append((px,py)); codes.append(L)

        if verts:
            patch = PathPatch(Path(verts, codes), fill=False,
                              linewidth=3, edgecolor='black')
            self.add_patch(patch)

            ## self.add_artist(Line2D(
            ##     [x,px], [y,py], lw=3, solid_capstyle="butt", color="black"
            ##     ))

    def hl(self, s):
        nodes = self.root.findall(s)
        if nodes:
            self.highlight(nodes)

    def hlines(self, nodes, width=5, color="red", xoff=0, yoff=0):
        offset = IdentityTransform()
        segs = []; w = []; o = []
        for n in filter(lambda x:x.parent, nodes):
            c = self.n2c[n]; p = self.n2c[n.parent]
            segs.append(((p.x,c.y),(c.x,c.y)))
            w.append(width); o.append((xoff,yoff))
        lc = LineCollection(segs, linewidths=w, transOffset=offset, offsets=o)
        lc.set_color(color)
        Axes.add_collection(self, lc)
        ## self.drawstack.append(("hlines", [nodes], dict(width=width,
        ##                                                color=color,
        ##                                                xoff=xoff,
        ##                                                yoff=yoff)))
        self.figure.canvas.draw_idle()
        return lc

    def hardcopy(self, relwidth=0.5, leafpad=1.5):
        p = HC.TreeFigure(self.root, relwidth=relwidth, leafpad=leafpad,
                          name=self.name, support=self.support,
                          leaf_fontsize=self.leaf_fontsize,
                          branch_fontsize=self.branch_fontsize,
                          branch_width=self.branch_width,
                          branch_color=self.branch_color,
                          highlight_support=self.highlight_support,
                          branchlabels=self.branchlabels,
                          decorators=self.decorators,
                          leaflabels=self.leaflabels,
                          mark_named=self._mark_named,
                          xlim=self.get_xlim(),
                          ylim=self.get_ylim())
        return p

    def highlight(self, nodes=None, width=5, color="red"):
        if self.highlightpatch:
            try:
                self.highlightpatch.remove()
            except:
                pass
        if not nodes:
            return

        if len(nodes)>1:
            mrca = self.root.mrca(nodes)
            if not mrca:
                return
        else:
            mrca = list(nodes)[0]

        M = Path.MOVETO; L = Path.LINETO
        verts = []
        codes = []
        seen = set()
        for node, coords in [ x for x in self.n2c.items() if x[0] in nodes ]:
            x = coords.x; y = coords.y
            p = node.parent
            while p:
                pcoords = self.n2c[p]
                px = pcoords.x; py = pcoords.y
                if node not in seen:
                    verts.append((x, y)); codes.append(M)
                    verts.append((px, y)); codes.append(L)
                    verts.append((px, py)); codes.append(L)
                    seen.add(node)
                if p == mrca or node == mrca:
                    break
                node = p
                coords = self.n2c[node]
                x = coords.x; y = coords.y
                p = node.parent
        px, py = verts[-1]
        verts.append((px, py)); codes.append(M)

        self.highlightpath = Path(verts, codes)
        self.highlightpatch = PathPatch(
            self.highlightpath, fill=False, linewidth=width, edgecolor=color,
            capstyle='round', joinstyle='round'
            )
        return self.add_patch(self.highlightpatch)

    def find(self, s):
        """
        Find node(s) matching pattern s and zoom to node(s)
        """
        nodes = list(self.root.find(s))
        if nodes:
            self.zoom_nodes(nodes)

    def zoom_nodes(self, nodes, border=1.2):
        y0, y1 = self.get_ylim(); x0, x1 = self.get_xlim()
        y0 = max(0, y0); y1 = min(1, y1)

        n2c = self.n2c
        v = [ n2c[n] for n in nodes ]
        ymin = min([ c.y for c in v ])
        ymax = max([ c.y for c in v ])
        xmin = min([ c.x for c in v ])
        xmax = max([ c.x for c in v ])
        bb = Bbox(((xmin,ymin), (xmax, ymax)))

        # convert data coordinates to display coordinates
        transform = self.transData.transform
        disp_bb = [Bbox(transform(bb))]
        for n in nodes:
            if n.isleaf:
                txt = self.node2label[n]
                if txt.get_visible():
                    disp_bb.append(txt.get_window_extent())

        disp_bb = Bbox.union(disp_bb).expanded(border, border)

        # convert back to data coordinates
        points = self.transData.inverted().transform(disp_bb)
        x0, x1 = points[:,0]
        y0, y1 = points[:,1]
        self.set_xlim(x0, x1)
        self.set_ylim(y0, y1)

    def zoom_clade(self, anc, border=1.2):
        if anc.isleaf:
            self.center_node(anc)

        else:
            self.zoom_nodes(list(anc), border)

    def draw_leaf_labels(self, *args):
        leaves = list(filter(lambda x:x[0].isleaf,
                             self.get_visible_nodes(labeled_only=True)))
        psep = self.leaf_pixelsep()
        fontsize = min(self.leaf_fontsize, max(psep, 8))
        n2l = self.node2label
        transform = self.transData.transform
        sub = operator.sub

        for n in leaves:
            n2l[n[0]].set_visible(False)

        # draw leaves
        leaves_drawn = []
        for n, x, y in leaves:
            txt = self.node2label[n]
            if not leaves_drawn:
                txt.set_visible(True)
                leaves_drawn.append(txt)
                self.figure.canvas.draw_idle()
                continue

            txt2 = leaves_drawn[-1]
            y0 = y; y1 = txt2.xy[1]
            sep = sub(*transform(([0,y0],[0,y1]))[:,1])
            if sep > fontsize:
                txt.set_visible(True)
                txt.set_size(fontsize)
                leaves_drawn.append(txt)
        self.figure.canvas.draw_idle()

        if leaves_drawn:
            leaves_drawn[0].set_size(fontsize)

        return fontsize

    def draw_labels(self, *args):
        fs = max(10, self.draw_leaf_labels())
        nodes = self.get_visible_nodes(labeled_only=True)
        ## print [ x[0].id for x in nodes ]
        branches = list(filter(lambda x:(not x[0].isleaf), nodes))
        n2l = self.node2label
        for n, x, y in branches:
            t = n2l[n]
            t.set_visible(True)
            t.set_size(fs)

    def unclutter(self, *args):
        nodes = self.get_visible_nodes(labeled_only=True)
        branches = list(filter(lambda x:(not x[0].isleaf), nodes))
        psep = self.leaf_pixelsep()
        n2l = self.node2label
        fontsize = min(self.leaf_fontsize*1.2, max(psep, self.leaf_fontsize))

        drawn = []
        for n, x, y in branches:
            txt = n2l[n]
            try:
                bb = txt.get_window_extent().expanded(2, 2)
                vis = True
                for n2 in reversed(drawn):
                    txt2 = n2l[n2]
                    if bb.overlaps(txt2.get_window_extent()):
                        txt.set_visible(False)
                        vis = False
                        self.figure.canvas.draw_idle()
                        break
                if vis:
                    txt.set_visible(True)
                    txt.set_size(fontsize)
                    self.figure.canvas.draw_idle()
                    drawn.append(n)
            except RuntimeError:
                pass
                ## txt.set_visible(True)
                ## txt.set_size(fontsize)
                ## drawn.append(n)
                ## self.figure.canvas.draw_idle()

    def leaf_pixelsep(self):
        y0, y1 = self.get_ylim()
        y0 = max(0, y0)
        y1 = min(self.nleaves, y1)
        display_points = self.transData.transform(((0, y0), (0, y1)))
        # height in pixels (visible y data extent)
        height = operator.sub(*reversed(display_points[:,1]))
        pixelsep = height/((y1-y0)/self.leaf_hsep)
        return pixelsep

    def ypp(self):
        y0, y1 = self.get_ylim()
        p0, p1 = self.transData.transform(((0, y0), (0, y1)))[:,1]
        return (y1-y0)/float(p1-p0)

    def draw_labels_old(self, *args):
        if self.nleaves:
            y0, y1 = self.get_ylim()
            y0 = max(0, y0); y1 = min(1, y1)

            display_points = self.transData.transform(((0, y0), (0, y1)))
            # height in pixels (visible y data extent)
            height = operator.sub(*reversed(display_points[:,1]))
            pixelsep = height/((y1-y0)/self.leaf_hsep)
            fontsize = min(max(pixelsep-2, 8), 12)

            if pixelsep >= 8:
                for node, txt in self.node2label.items():
                    if node.isleaf:
                        if self.leaflabels:
                            c = self.n2c[node]
                            x = c.x; y = c.y
                            if (y0 < y < y1):
                                txt.set_size(fontsize)
                                txt.set_visible(True)
                    else:
                        if self.branchlabels:
                            c = self.n2c[node]
                            x = c.x; y = c.y
                            if (y0 < y < y1):
                                txt.set_size(fontsize)
                                txt.set_visible(True)
            elif pixelsep >= 4:
                for node, txt in self.node2label.items():
                    if node.isleaf:
                        txt.set_visible(False)
                    else:
                        if self.branchlabels:
                            c = self.n2c[node]
                            x = c.x; y = c.y
                            if (y0 < y < y1):
                                txt.set_size(fontsize)
                                txt.set_visible(True)
            else:
                for node, txt in self.node2label.items():
                    txt.set_visible(False)
            self.figure.canvas.draw_idle()

    def redraw(self, home=False, layout=True):
        """
        Replot the tree
        """
        xlim = self.get_xlim()
        ylim = self.get_ylim()
        self.cla()
        if layout:
            self.layout()
        self.plot_tree()
        if self.interactive:
            self.callbacks.connect("ylim_changed", self.draw_labels)

        if home:
            self.home()
        else:
            self.set_xlim(*xlim)
            self.set_ylim(*ylim)

    def set_name(self, name):
        self.name = name
        if name:
            at = AnchoredText(
                self.name, loc=2, frameon=True,
                prop=dict(size=12, weight="bold")
                )
            at.patch.set_linewidth(0)
            at.patch.set_facecolor("white")
            at.patch.set_alpha(0.6)
            self.add_artist(at)
            return at

    def _path_to_parent(self, node):
        """
        For use in drawing branches
        """
        c = self.n2c[node]; x = c.x; y = c.y
        pc = self.n2c[node.parent]; px = pc.x; py = pc.y
        M = Path.MOVETO; L = Path.LINETO
        verts = [(x, y), (px, y), (px, py)]
        codes = [M, L, L]
        return verts, codes
        ## return [PathPatch(Path(verts, codes), fill=False,
        ##                   linewidth=width or self.branch_width,
        ##                   edgecolor=color or self.branch_color)]

    def layout(self):
        self.n2c = cartesian(self.root, scaled=self.scaled, yunit=1.0,
                             smooth=self.smooth_xpos)
        for c in self.n2c.values():
            c.x += self.xoff; c.y += self.yoff
        sv = sorted([
            [c.y, c.x, n] for n, c in self.n2c.items()
            ])
        self.coords = sv#numpy.array(sv)
        ## n2c = self.n2c
        ## self.node2linesegs = {}
        ## for node, coords in n2c.items():
        ##     x = coords.x; y = coords.y
        ##     v = [(x,y)]
        ##     if node.parent:
        ##         pcoords = n2c[node.parent]
        ##         px = pcoords.x; py = pcoords.y
        ##         v.append((px,y))
        ##         v.append((px,py))
        ##         self.node2linesegs[node] = v

    def set_root(self, root):
        self.root = root
        self.leaves = root.leaves()
        self.nleaves = len(self.leaves)
        self.leaf_hsep = 1.0#/float(self.nleaves)

        for n in root.descendants():
            if n.length is None:
                self.scaled=False; break
        self.layout()

    def plot_tree(self, root=None, **kwargs):
        """
        Draw branches and labels
        """
        if root and not self.root:
            self.set_root(root)

        if self.interactive: pyplot.ioff()

        if "branchlabels" in kwargs:
            self.branchlabels = kwargs["branchlabels"]
        if "leaflabels" in kwargs:
            self.leaflabels = kwargs["leaflabels"]
        self.yaxis.set_visible(False)
        self.create_branch_artists()
        self.create_label_artists()
        if self.highlight_support:
            self.do_highlight_support()
        self.mark_named()
        ## self.home()

        for k, v in self.decorators:
            func, args, kwargs = v
            func(self, *args, **kwargs)

        self.set_name(self.name)
        self.adjust_xspine()

        if self.interactive: pyplot.ion()

        labels = [ x.label for x in self.root.leaves() ]
        def fmt(x, pos=None):
            if x<0: return ""
            try: return labels[int(round(x))]
            except: pass
            return ""
        #self.yaxis.set_major_formatter(FuncFormatter(fmt))

        return self

    def clade_dimensions(self):
        n2c = self.n2c
        d = {}
        def recurse(n, n2c, d):
            v = []
            for c in n.children:
                recurse(c, n2c, d)
                if c.isleaf:
                    x, y = n2c[c].point()
                    x0 = x1 = x; y0 = y1 = y
                else:
                    x0, x1, y0, y1 = d[c]
                v.append((x0, x1, y0, y1))
            if v:
                x0 = n2c[n].x
                x1 = max([ x[1] for x in v ])
                y0 = min([ x[2] for x in v ])
                y1 = max([ x[3] for x in v ])
                d[n] = (x0, x1, y0, y1)
        recurse(self.root, n2c, d)
        return d

    def clade_height_pixels(self):
        ypp = self.ypp()
        d = self.clade_dimensions()
        h = {}
        for n, (x0, x1, y0, y1) in d.items():
            h[n] = (y1-y0)/ypp
        return h

    def _decimate_nodes(self, n=500):
        leaves = self.leaves
        nleaves = len(leaves)
        if nleaves > n:
            indices = numpy.linspace(0, nleaves-1, n).astype(int)
            leaves = [ leaves[i] for i in indices ]
            return set(list(chain.from_iterable([ list(x.rootpath())
                                                  for x in leaves ])))
        else:
            return self.root

    def create_branch_artists(self):
        """
        Use MPL Paths to draw branches
        """
        ## patches = []
        verts = []; codes = []
        for node in self.root.descendants():
            v, c = self._path_to_parent(node)
            verts.extend(v); codes.extend(c)
        self.branchpatch = PathPatch(
            Path(verts, codes), fill=False,
            linewidth=self.branch_width,
            edgecolor=self.branch_color
            )
        self.add_patch(self.branchpatch)
        ## for node in self._decimate_nodes():
        ##     if node.parent:
        ##         for p in self._path_to_parent(node):
        ##             patches.append(p)
        ## self.branch_patches = PatchCollection(patches, match_original=True)
        ## self.add_collection(self.branch_patches)

        ## print "enter: create_branch_artists"
        ## self.node2branch = {}
        ## for node, segs in self.node2linesegs.items():
        ##     line = Line2D(
        ##         [x[0] for x in segs], [x[1] for x in segs],
        ##         lw=self.branch_width, color=self.branch_color
        ##         )
        ##     line.set_visible(False)
        ##     Axes.add_artist(self, line)
        ##     self.node2branch[node] = line

        ## d = self.node2linesegs
        ## segs = [ d[n] for n in self.root if (n in d) ]

        ## dims = self.clade_dimensions(); ypp = self.ypp()
        ## def recurse(n, dims, clades, terminals):
        ##     stop = False
        ##     h = None
        ##     v = dims.get(n)
        ##     if v: h = (v[3]-v[2])/ypp
        ##     if (h and (h < 20)) or (not h):
        ##         stop = True
        ##         terminals.append(n)
        ##     if not stop:
        ##         clades.append(n)
        ##         for c in n.children:
        ##             recurse(c, dims, clades, terminals)
        ## clades = []; terminals = []
        ## recurse(self.root, dims, clades, terminals)
        ## segs = [ d[n] for n in self.root if (n in d) and (n in clades) ]
        ## for t in terminals:
        ##     if t.isleaf:
        ##         segs.append(d[t])
        ##     else:
        ##         x0, x1, y0, y1 = dims[t]
        ##         x, y = self.n2c[t].point()
        ##         px, py = self.n2c[t.parent].point()
        ##         segs.append(((px,py), (px,y), (x,y), (x1, y0), (x1,y1), (x,y)))

        ## lc = LineCollection(segs, linewidths=self.branch_width,
        ##                     colors = self.branch_color)
        ## self.branches_linecollection = Axes.add_collection(self, lc)
        ## print "leave: create_branch_artists"

    def create_label_artists(self):
        ## print "enter: create_label_artists"
        self.node2label = {}
        n2c = self.n2c
        for node, coords in n2c.items():
            x = coords.x; y = coords.y
            if node.isleaf and node.label and self.leaflabels:
                txt = self.annotate(
                    node.label,
                    xy=(x, y),
                    xytext=(self.leaf_offset, 0),
                    textcoords="offset points",
                    verticalalignment=self.leaf_valign,
                    horizontalalignment=self.leaf_halign,
                    fontsize=self.leaf_fontsize,
                    clip_on=True,
                    picker=True
                )
                txt.node = node
                txt.set_visible(False)
                self.node2label[node] = txt

            if (not node.isleaf) and node.label and self.branchlabels:
                txt = self.annotate(
                    node.label,
                    xy=(x, y),
                    xytext=(self.branch_offset,0),
                    textcoords="offset points",
                    verticalalignment=self.branch_valign,
                    horizontalalignment=self.branch_halign,
                    fontsize=self.branch_fontsize,
                    bbox=dict(fc="lightyellow", ec="none", alpha=0.8),
                    clip_on=True,
                    picker=True
                )
                ## txt.set_visible(False)
                txt.node = node
                self.node2label[node] = txt
        ## print "leave: create_label_artists"

    def adjust_xspine(self):
        v = sorted([ c.x for c in self.n2c.values() ])
        try:
            self.spines["bottom"].set_bounds(v[0],v[-1])
        except AttributeError:
            pass
        for t,n,s in self.xaxis.iter_ticks():
            if (n > v[-1]) or (n < v[0]):
                t.set_visible(False)

    def mark_named(self):
        if self._mark_named:
            n2c = self.n2c
            cv = [ c for n, c in n2c.items() if n.label and (not n.isleaf) ]
            x = [ c.x for c in cv ]
            y = [ c.y for c in cv ]
            if x and y:
                self.scatter(x, y, s=5, color='black')

    def home(self):
        td = self.transData
        trans = td.inverted().transform
        xmax = xmin = ymax = ymin = 0
        if self.node2label:
            try:
                v = [ x.get_window_extent() for x in self.node2label.values()
                      if x.get_visible() ]
                if v:
                    xmax = trans((max([ x.xmax for x in v ]),0))[0]
                    xmin = trans((min([ x.xmin for x in v ]),0))[0]
            except RuntimeError:
                pass

        v = self.n2c.values()
        ymin = min([ c.y for c in v ])
        ymax = max([ c.y for c in v ])
        xmin = min(xmin, min([ c.x for c in v ]))
        xmax = max(xmax, max([ c.x for c in v ]))
        xspan = xmax - xmin; xpad = xspan*0.05
        yspan = ymax - ymin; ypad = yspan*0.05
        self.set_xlim(xmin-xpad, xmax+xpad*2)
        self.set_ylim(ymin-ypad, ymax+ypad)
        self.adjust_xspine()

    def scroll(self, x, y):
        x0, x1 = self.get_xlim()
        y0, y1 = self.get_ylim()
        xd = (x1-x0)*x
        yd = (y1-y0)*y
        self.set_xlim(x0+xd, x1+xd)
        self.set_ylim(y0+yd, y1+yd)
        self.adjust_xspine()

    def plot_labelcolor(self, nodemap, state2color=None):
        if state2color is None:
            c = colors.tango()
            states = sorted(set(nodemap.values()))
            state2color = dict(zip(states, c))

        for node, txt in self.node2label.items():
            s = nodemap.get(node)
            if s is not None:
                c = state2color[s]
                if c:
                    txt.set_color(c)
        self.figure.canvas.draw_idle()

    def node_image(self, node, imgfile, maxdim=100, border=0):
        xoff = self.leaf_offset
        n = self.root[node]; c = self.n2c[n]; p = (c.x, c.y)
        img = Image.open(imgfile)
        if max(img.size) > maxdim:
            img.thumbnail((maxdim, maxdim))
        imgbox = OffsetImage(img)
        xycoords = self.node2label.get(node) or "data"
        if xycoords != "data": p = (1, 0.5)
        abox = AnnotationBbox(imgbox, p,
                              xybox=(xoff, 0.0),
                              xycoords=xycoords,
                              box_alignment=(0.0,0.5),
                              pad=0.0,
                              boxcoords=("offset points"))
        self.add_artist(abox)

    def plot_discrete(self, data, cmap=None, name=None,
                      xoff=10, yoff=0, size=15, legend=1):
        root = self.root
        if cmap is None:
            import ivy
            c = colors.tango()
            states = sorted(set(data.values()))
            cmap = dict(zip(states, c))
        n2c = self.n2c
        points = []; c = []
        d = dict([ (n, data.get(n)) for n in root if data.get(n) is not None ])
        for n, v in d.items():
            coord = n2c[n]
            points.append((coord.x, coord.y)); c.append(cmap[v])

        boxes = symbols.squares(self, points, c, size, xoff=xoff, yoff=yoff)

        if legend:
            handles = []; labels = []
            for v, c in sorted(cmap.items()):
                handles.append(Rectangle((0,0),0.5,1,fc=c))
                labels.append(str(v))
            self.legend(handles, labels, loc=legend)

        self.figure.canvas.draw_idle()
        return boxes

    def plot_continuous(self, data, mid=None, name=None, cmap=None,
                        size=15, colorbar=True):
        area = (size*0.5)*(size*0.5)*numpy.pi
        values = data.values()
        vmin = min(values); vmax = max(values)
        if mid is None:
            mid = (vmin+vmax)*0.5
            delta = vmax-vmin*0.5
        else:
            delta = max(abs(vmax-mid), abs(vmin-mid))
        norm = mpl_colors.Normalize(mid-delta, mid+delta)
        ## if cmap is None: cmap = mpl_colormap.binary
        if cmap is None: cmap = mpl_colormap.hot
        n2c = self.n2c
        X = numpy.array(
            [ (n2c[n].x, n2c[n].y, v) for n, v in data.items() if n in n2c ]
            )
        circles = self.scatter(
            X[:,0], X[:,1], s=area, c=X[:,2], cmap=cmap, norm=norm,
            zorder=1000
            )
        if colorbar:
            cbar = self.figure.colorbar(circles, ax=self, shrink=0.7)
            if name:
                cbar.ax.set_xlabel(name)

        self.figure.canvas.draw_idle()

class RadialTree(Tree):
    def layout(self):
        from ..layout_polar import calc_node_positions
        start = self.start if hasattr(self, 'start') else 0
        end = self.end if hasattr(self, 'end') else None
        self.n2c = calc_node_positions(self.root, scaled=self.scaled,
                                       start=start, end=end)
        sv = sorted([
            [c.y, c.x, n] for n, c in self.n2c.items()
            ])
        self.coords = sv

    ## def _path_to_parent(self, node, width=None, color=None):
    ##     c = self.n2c[node]; theta1 = c.angle; r = c.depth
    ##     M = Path.MOVETO; L = Path.LINETO
    ##     pc = self.n2c[node.parent]; theta2 = pc.angle
    ##     px1 = math.cos(math.radians(c.angle))*pc.depth
    ##     py1 = math.sin(math.radians(c.angle))*pc.depth
    ##     verts = [(c.x,c.y),(px1,py1)]; codes = [M,L]
    ##     #verts.append((pc.x,pc.y)); codes.append(L)
    ##     path = PathPatch(Path(verts, codes), fill=False,
    ##                      linewidth=width or self.branch_width,
    ##                      edgecolor=color or self.branch_color)
    ##     diam = pc.depth*2
    ##     t1, t2 = tuple(sorted((theta1,theta2)))
    ##     arc = Arc((0,0), diam, diam, theta1=t1, theta2=t2,
    ##               edgecolor=color or self.branch_color,
    ##               linewidth=width or self.branch_width)
    ##     return [path, arc]

    def _path_to_parent(self, node):
        c = self.n2c[node]; theta1 = c.angle; r = c.depth
        M = Path.MOVETO; L = Path.LINETO
        pc = self.n2c[node.parent]; theta2 = pc.angle
        px1 = math.cos(math.radians(c.angle))*pc.depth
        py1 = math.sin(math.radians(c.angle))*pc.depth
        verts = [(c.x,c.y),(px1,py1)]; codes = [M,L]
        t1, t2 = tuple(sorted((theta1,theta2)))
        diam = pc.depth*2
        arc = Arc((0,0), diam, diam, theta1=t1, theta2=t2)
        arcpath = arc.get_path()
        av = arcpath.vertices * pc.depth
        ac = arcpath.codes
        verts.extend(av.tolist())
        codes.extend(ac.tolist())
        return verts, codes

    def highlight(self, nodes=None, width=5, color="red"):
        if self.highlightpatch:
            try:
                self.highlightpatch.remove()
            except:
                pass
        if not nodes:
            return

        if len(nodes)>1:
            mrca = self.root.mrca(nodes)
            if not mrca:
                return
        else:
            mrca = list(nodes)[0]

        M = Path.MOVETO; L = Path.LINETO
        verts = []
        codes = []
        seen = set()
        patches = []
        for node, coords in [ x for x in self.n2c.items() if x[0] in nodes ]:
            x = coords.x; y = coords.y
            p = node.parent
            while p:
                pcoords = self.n2c[p]
                px = pcoords.x; py = pcoords.y
                if node not in seen:
                    v, c = self._path_to_parent(node)
                    verts.extend(v)
                    codes.extend(c)
                    seen.add(node)
                if p == mrca or node == mrca:
                    break
                node = p
                coords = self.n2c[node]
                x = coords.x; y = coords.y
                p = node.parent
        ## px, py = verts[-1]
        ## verts.append((px, py)); codes.append(M)
        self.highlightpath = Path(verts, codes)
        self.highlightpatch = PathPatch(
            self.highlightpath, fill=False, linewidth=width, edgecolor=color
            )
        self.add_patch(self.highlightpatch)
        ## self.highlight_patches = PatchCollection(patches, match_original=True)
        ## self.add_collection(self.highlight_patches)

class OverviewTree(Tree):
    def __init__(self, *args, **kwargs):
        kwargs["leaflabels"] = False
        kwargs["branchlabels"] = False
        Tree.__init__(self, *args, **kwargs)
        self.xaxis.set_visible(False)
        self.spines["bottom"].set_visible(False)
        self.add_overview_rect()

    def set_target(self, target):
        self.target = target

    def add_overview_rect(self):
        rect = UpdatingRect([0, 0], 0, 0, facecolor='black', edgecolor='red')
        rect.set_alpha(0.2)
        rect.target = self.target
        rect.set_bounds(*self.target.viewLim.bounds)
        self.zoomrect = rect
        self.add_patch(rect)
        ## if pyplot.isinteractive():
        self.target.callbacks.connect('xlim_changed', rect)
        self.target.callbacks.connect('ylim_changed', rect)

    def redraw(self):
        Tree.redraw(self)
        self.add_overview_rect()
        self.figure.canvas.draw_idle()

def axes_enter(e):
    ax = e.inaxes
    ax._active = True

def axes_leave(e):
    ax = e.inaxes
    ax._active = False

def onselect(estart, estop):
    b = estart.button
    ## print b, estart.key

def onkeypress(e):
    ax = e.inaxes
    k = e.key
    if ax and k == 't':
        ax.home()
    if ax and k == "down":
        ax.scroll(0, -0.1)
        ax.figure.canvas.draw_idle()
    if ax and k == "up":
        ax.scroll(0, 0.1)
        ax.figure.canvas.draw_idle()
    if ax and k == "left":
        ax.scroll(-0.1, 0)
        ax.figure.canvas.draw_idle()
    if ax and k == "right":
        ax.scroll(0.1, 0)
        ax.figure.canvas.draw_idle()
    if ax and k and k in '=+':
        ax.zoom(0.1,0.1)
    if ax and k == '-':
        ax.zoom(-0.1,-0.1)

def ondrag(e):
    ax = e.inaxes
    button = e.button
    if ax and button == 2:
        if not ax.pan_start:
            ax.pan_start = (e.xdata, e.ydata)
            return
        x, y = ax.pan_start
        xdelta = x - e.xdata
        ydelta = y - e.ydata
        x0, x1 = ax.get_xlim()
        xspan = x1-x0
        y0, y1 = ax.get_ylim()
        yspan = y1 - y0
        midx = (x1+x0)*0.5
        midy = (y1+y0)*0.5
        ax.set_xlim(midx+xdelta-xspan*0.5, midx+xdelta+xspan*0.5)
        ax.set_ylim(midy+ydelta-yspan*0.5, midy+ydelta+yspan*0.5)
        ax.adjust_xspine()

def onbuttonrelease(e):
    ax = e.inaxes
    button = e.button
    if button == 2:
        ## print "pan end"
        ax.pan_start = None

def onpick(e):
    ax = e.mouseevent.inaxes
    if ax:
        ax.picked(e)

def onscroll(e):
    ax = e.inaxes
    if ax:
        b = e.button
        ## print b
        k = e.key
        if k == None and b =="up":
            ax.zoom(0.1,0.1)
        if k == None and b =="down":
            ax.zoom(-0.1,-0.1)
        if k == "shift" and b == "up":
            ax.zoom_cxy(0.1, 0, e.xdata, e.ydata)
        if k == "shift" and b == "down":
            ax.zoom_cxy(-0.1, 0, e.xdata, e.ydata)
        if k == "control" and b == "up":
            ax.zoom_cxy(0, 0.1, e.xdata, e.ydata)
        if k == "control" and b == "down":
            ax.zoom_cxy(0, -0.1, e.xdata, e.ydata)
        if k == "d" and b == "up":
            ax.scroll(0, 0.1)
        if (k == "d" and b == "down"):
            ax.scroll(0, -0.1)
        if k == "c" and b == "up":
            ax.scroll(-0.1, 0)
        if k == "c" and b == "down":
            ax.scroll(0.1, 0)
        try: ax.adjust_xspine()
        except: pass
        ax.figure.canvas.draw_idle()

def onclick(e):
    ax = e.inaxes
    if ax and e.button==1 and hasattr(ax, "zoomrect") and ax.zoomrect:
        # overview clicked; reposition zoomrect
        r = ax.zoomrect
        x = e.xdata
        y = e.ydata
        arr = ax.transData.inverted().transform(r.get_extents())
        xoff = (arr[1][0]-arr[0][0])*0.5
        yoff = (arr[1][1]-arr[0][1])*0.5
        r.target.set_xlim(x-xoff,x+xoff)
        r.target.set_ylim(y-yoff,y+yoff)
        r(r.target)
        ax.figure.canvas.draw_idle()

    if ax and e.button==2:
        ## print "pan start", (e.xdata, e.ydata)
        ax.pan_start = (e.xdata, e.ydata)

def test_decorate(treeplot):
    import evolve
    data = evolve.brownian(treeplot.root)
    values = data.values()
    vmin = min(values); vmax = max(values)
    norm = mpl_colors.Normalize(vmin, vmax)
    cmap = mpl_colormap.binary
    n2c = treeplot.n2c
    X = numpy.array(
        [ (n2c[n].x, n2c[n].y, v)
          for n, v in data.items() if n in n2c ]
        )
    circles = treeplot.scatter(
        X[:,0], X[:,1], s=200, c=X[:,2], cmap=cmap, norm=norm,
        zorder=100
        )

class Decorator(object):
    def __init__(self, treeplot):
        self.plot = treeplot

class VisToggle(object):
    def __init__(self, name, treeplot=None, value=False):
        self.name = name
        self.plot = treeplot
        self.value = value

    def __nonzero__(self):
        return self.value

    def __repr__(self):
        return "%s: %s" % (self.name, self.value)

    def redraw(self):
        if self.plot:
            self.plot.redraw()

    def toggle(self):
        self.value = not self.value
        self.redraw()

    def show(self):
        if self.value == False:
            self.value = True
            self.redraw()

    def hide(self):
        if self.value == True:
            self.value = False
            self.redraw()

TreePlot = subplot_class_factory(Tree)
RadialTreePlot = subplot_class_factory(RadialTree)
OverviewTreePlot = subplot_class_factory(OverviewTree)

if __name__ == "__main__":
    import evolve
    root, data = evolve.test_brownian()
    plot_continuous(root, data, name="Brownian", mid=0.0)