#! /usr/bin/env python

##############################################################################
##  DendroPy Phylogenetic Computing Library.
##
##  Copyright 2010-2015 Jeet Sukumaran and Mark T. Holder.
##  All rights reserved.
##
##  See "LICENSE.rst" for terms and conditions of usage.
##
##  If you use this work or any portion thereof in published work,
##  please cite it as:
##
##     Sukumaran, J. and M. T. Holder. 2010. DendroPy: a Python library
##     for phylogenetic computing. Bioinformatics 26: 1569-1571.
##
##############################################################################

"""
Functions, classes, and methods for working with Kingman's n-coalescent
framework.
"""

import math
import dendropy
from dendropy.utility import GLOBAL_RNG
from dendropy.utility import constants
from dendropy.calculate import probability
from dendropy.calculate import combinatorics

###############################################################################
## Calculations and statistics

def discrete_time_to_coalescence(n_genes,
                                 pop_size=None,
                                 n_to_coalesce=2,
                                 rng=None):
    """
    A random draw from the "Kingman distribution" (discrete time version): Time
    to go from ``n_genes`` genes to ``n_genes``-1 genes in a discrete-time
    Wright-Fisher population of ``pop_size`` genes; i.e. waiting time until
    ``n-genes`` lineages coalesce in a population of ``pop_size`` genes.

    Parameters
    ----------

    n_genes : integer
        The number of genes in the sample.
    pop_size : integer
        The effective *haploid* population size; i.e., number of genes in the
        population: 2 * N in a diploid population of N individuals, or N in a
        haploid population of N individuals.
    n_to_coalesce : integer
        The waiting time that will be returned will be the waiting time for
        this number of genes in the sample to coalesce.
    rng : ``Random`` object
        The random number generator instance.

    Returns
    -------
    k : integer
        A randomly-generated waiting time (in discrete generations) for
        ``n_to_coalesce`` genes to coalesce out of a sample of ``n_genes`` in a
        population of ``pop_size`` genes.

    """
    if not pop_size:
        time_units = 1.0
    else:
        time_units = pop_size
    if rng is None:
        rng = GLOBAL_RNG
    p = pop_size / combinatorics.choose(n_genes, n_to_coalesce)
    tmrca = probability.geometric_rv(p)
    return tmrca * time_units

def time_to_coalescence(n_genes,
        pop_size=None,
        n_to_coalesce=2,
        rng=None):
    """
    A random draw from the "Kingman distribution" (discrete time version): Time
    to go from ``n_genes`` genes to ``n_genes``-1 genes in a continuous-time
    Wright-Fisher population of ``pop_size`` genes; i.e. waiting time until
    ``n-genes`` lineages coalesce in a population of ``pop_size`` genes.

    Given the number of gene lineages in a sample, ``n_genes``, and a
    population size, ``pop_size``, this function returns a random number from
    an exponential distribution with rate $\\choose(``pop_size``, 2)$.
    ``pop_size`` is the effective *haploid* population size; i.e., number of gene
    in the population: 2 * N in a diploid population of N individuals,
    or N in a haploid population of N individuals. If ``pop_size`` is 1 or 0 or
    None, then time is in haploid population units; i.e. where 1 unit of time
    equals 2N generations for a diploid population of size N, or N generations
    for a haploid population of size N. Otherwise time is in generations.

    The coalescence time, or the waiting time for the coalescence, of two
    gene lineages evolving in a population with haploid size $N$ is an
    exponentially-distributed random variable with rate of $N$ an
    expectation of $\\frac{1}{N}$).
    The waiting time for coalescence of *any* two gene lineages in a sample of
    $n$ gene lineages evolving in a population with haploid size $N$ is an
    exponentially-distributed random variable with rate of $\\choose{N, 2}$ and
    an expectation of $\\frac{1}{\choose{N, 2}}$.

    Parameters
    ----------
    n_genes : integer
        The number of genes in the sample.
    pop_size : integer
        The effective *haploid* population size; i.e., number of genes in the
        population: 2 * N in a diploid population of N individuals, or N in a
        haploid population of N individuals.
    n_to_coalesce : integer
        The waiting time that will be returned will be the waiting time for
        this number of genes in the sample to coalesce.
    rng : ``Random`` object
        The random number generator instance to use.

    Returns
    -------
    k : float
        A randomly-generated waiting time (in continuous time) for
        ``n_to_coalesce`` genes to coalesce out of a sample of ``n_genes`` in a
        population of ``pop_size`` genes.
    """
    if rng is None:
        rng = GLOBAL_RNG
    if not pop_size:
        time_units = 1.0
    else:
        time_units = pop_size
    rate = combinatorics.choose(n_genes, n_to_coalesce)
    tmrca = rng.expovariate(rate)
    return tmrca * time_units

def expected_tmrca(n_genes, pop_size=None, n_to_coalesce=2):
    """
    Expected (mean) value for the Time to the Most Recent Common Ancestor of
    ``n_to_coalesce`` genes in a sample of ``n_genes`` drawn from a population of
    ``pop_size`` genes.

    Parameters
    ----------
    n_genes : integer
        The number of genes in the sample.
    pop_size : integer
        The effective *haploid* population size; i.e., number of genes in the
        population: 2 * N in a diploid population of N individuals, or N in a
        haploid population of N individuals.
    n_to_coalesce : integer
        The waiting time that will be returned will be the waiting time for
        this number of genes in the sample to coalesce.
    rng : ``Random`` object
        The random number generator instance.

    Returns
    -------
    k : float
        The expected waiting time (in continuous time) for ``n_to_coalesce``
        genes to coalesce out of a sample of ``n_genes`` in a population of
        ``pop_size`` genes.

    """
    nc2 = combinatorics.choose(n_genes, n_to_coalesce)
    tmrca = (float(1)/nc2)
    if pop_size is not None:
        return tmrca * pop_size
    else:
        return tmrca

def coalesce_nodes(nodes,
             pop_size=None,
             period=None,
             rng=None,
             use_expected_tmrca=False):
    """
    Returns a list of nodes that have not yet coalesced once ``period`` is
    exhausted.

    This function will a draw a coalescence time, ``t``, from an exponential
    distribution with a rate of ``choose(k, 2)``, where ``k`` is the number of
    nodes. If ``period`` is given and if this time is less than ``period``, or if
    ``period`` is not given, then two nodes are selected at random from ``nodes``,
    and coalesced: a new node is created, and the two nodes are added as
    child_nodes to this node with an edge length such the the total length from
    tip to the ancestral node is equal to the depth of the deepest child + ``t``.
    The two nodes are removed from the list of nodes, and the new node is added
    to it. ``t`` is then deducted from ``period``, and the process repeats.

    The function ends and returns the list of nodes once ``period`` is
    exhausted or if any draw of ``t`` exceeds ``period``, if ``period`` is
    given or when there is only one node left.

    As each coalescent event occurs, *all* nodes have their edges
    extended to the point of the coalescent event. In the case of
    constrained coalescence, all uncoalesced nodes have their edges
    extended to the end of the period (coalesced nodes have the edges
    fixed by the coalescent event in their ancestor).  Thus multiple
    calls to this method with the same set of nodes will gradually
    'grow' the edges, until all the the nodes coalesce. The edge
    lengths of the nodes passed to this method thus should not be
    modified or reset until the process is complete.

    Parameters
    ----------
    nodes : iterable[|Node|]
        An interable of |Node| objects representing a sample of neutral
        genes (some, all, or none of these nodes may have descendent nodes).
    pop_size : integer
        The effective *haploid* population size; i.e., number of genes in the
        population: 2 * N in a diploid population of N individuals, or N in a
        haploid population of N individuals.
    period : numeric
        The time that the genes have to coalesce. If ``pop_size`` is 1 or 0 or
        None, then time is in haploid population units; i.e. where 1 unit of
        time equals 2N generations for a diploid population of size N, or N
        generations for a haploid population of size N. Otherwise time is in
        generations.
    rng : ``Random`` object
        The random number generator instance to use. If not specified, the
        default RNG will be used.
    use_expected_tmrca : bool
        If |True|, then instead of random times, the *expected* times will be
        used.

    Returns
    -------
    nodes : iterable[|Node|]
        A list of nodes once ``period`` is exhausted or if any draw of ``t``
        exceeds ``period``, if ``period`` is given or when there is only one node
        left.
    """

    # idiot-check, because I can be an idiot
    if not nodes:
        return []

    # set the random number generator
    if rng is None:
        rng = GLOBAL_RNG

    # define the function needed to create new coalescence nodes
    new_node = nodes[0].__class__

    # make a shallow copy of the node list
    nodes = list(nodes)

    # start tracking the time remaining
    time_remaining = period

    # If there is no time constraint, we want to continue coalescing
    # until there is only one gene left in the pool. If there is a
    # time constraint, we continue as long as there is time remaining,
    # but we do not control for that here: it is automatically taken
    # care of when the time drawn for the next coalescent event
    # exceeds the time remaining, and triggers a break from the loop
    while len(nodes) > 1:

        if use_expected_tmrca:
            tmrca = expected_tmrca(len(nodes), pop_size=pop_size)
        else:
            # draw a time to coalesce: this will be an exponential random
            # variable with parameter (rate) of BINOMIAL[n_genes 2]
            # multiplied pop_size
            tmrca = time_to_coalescence(len(nodes), pop_size=pop_size, rng=rng)

        # if no time_remaining is given (i.e, we want to coalesce till
        # there is only one gene left) or, if we are working under the
        # constrained coalescent, if the time to the next coalescence
        # event is not longer than the time_remaining
        if time_remaining is None or tmrca <= time_remaining:

            # stretch out the edges of all the nodes to this time
            for node in nodes:
                if node.edge.length is None:
                    node.edge.length = 0.0
                node.edge.length = node.edge.length + tmrca

            # pick two nodes to coalesce at random
            to_coalesce = rng.sample(nodes, 2)

            # create the new ancestor of these nodes
            new_ancestor = new_node()

            # add the nodes as child nodes of the new node, their
            # common ancestor, and set the ancestor's edge length
            new_ancestor.add_child(to_coalesce[0])
            new_ancestor.add_child(to_coalesce[1])
            new_ancestor.edge.length = 0.0

            # remove the nodes that have coalesced from the pool of
            # nodes
            nodes.remove(to_coalesce[0])
            nodes.remove(to_coalesce[1])

            # add the ancestor to the pool of nodes
            nodes.append(new_ancestor)

            # adjust the time_remaining left to coalesce
            if time_remaining is not None:
                time_remaining = time_remaining - tmrca
        else:
            # the next coalescent event takes place after the period constraint
            break

    # adjust the edge lengths of all the nodes, so they are at the
    # correct height, with the edges 'lining up' at the end of
    # coalescent period
    if time_remaining is not None and time_remaining > 0:
        for node in nodes:
            if node.edge.length is None:
                node.edge.length = 0.0
            node.edge.length = node.edge.length + time_remaining

    # return the list of nodes that have not coalesced
    return nodes

def node_waiting_time_pairs(tree, ultrametricity_precision=constants.DEFAULT_ULTRAMETRICITY_PRECISION):
    """
    Returns a list of tuples of (nodes, coalescent interval time) on the tree.
    That is, each element in the list is tuple pair consisting of where: the
    first element of the pair is an internal node representing a coalescent
    event on the tree, and the second element of the pair is the time between
    this coalescence event and the earlier (more recent) one.

    Parameters
    ----------
    tree : |Tree|
        A tree instance.
    ultrametricity_precision : float
        When calculating the node ages, an error will be raised if the tree is
        not ultrametric. This error may be due to floating-point or numerical
        imprecision. You can set the precision of the ultrametricity validation
        by setting the ``ultrametricity_precision`` parameter. E.g., use
        ``ultrametricity_precision=0.01`` for a more relaxed precision, down to
        2 decimal places. Use ``ultrametricity_precision=False`` to disable
        checking of ultrametricity.

    Returns
    -------
    x : list of tuples (node, coalescent interval)
        Returns list of tuples of (node, coalescent interval [= time between
        last coalescent event and current node age])
    """
    tree.calc_node_ages(ultrametricity_precision=ultrametricity_precision)
    ages = [(n, n.age) for n in tree.internal_nodes()]
    ages.sort(key=lambda x: x[1])
    intervals = []
    intervals.append(ages[0])
    for i, d in enumerate(ages[1:]):
        nd = d[0]
        prev_nd = ages[i][0]
        intervals.append( (nd, nd.age - prev_nd.age) )
    return intervals

def extract_coalescent_frames(tree, ultrametricity_precision=constants.DEFAULT_ULTRAMETRICITY_PRECISION):
    """
    Returns a list of tuples describing the coalescent frames on the tree. That
    is, each element in the list is tuple pair consisting of where: the first
    element of the pair is the number of separate lineages remaining on the
    tree at coalescence event, and the second element of the pair is the time
    between this coalescence event and the earlier (more recent) one.

    Parameters
    ----------
    tree : |Tree|
        A tree instance.
    ultrametricity_precision : float
        When calculating the node ages, an error will be raised if the tree is
        not ultrametric. This error may be due to floating-point or numerical
        imprecision. You can set the precision of the ultrametricity validation
        by setting the ``ultrametricity_precision`` parameter. E.g., use
        ``ultrametricity_precision=0.01`` for a more relaxed precision, down to
        2 decimal places. Use ``ultrametricity_precision=False`` to disable
        checking of ultrametricity.

    Returns
    -------
    x : dict
        Returns dictionary, with key = number of alleles, and values = waiting
        time for coalescent for the given tree
    """
    nwti = node_waiting_time_pairs(tree, ultrametricity_precision=ultrametricity_precision)
#     num_genes = len(tree.taxon_namespace)
    num_genes = len(tree.leaf_nodes())
    num_genes_wt = {}
    for n in nwti:
        num_genes_wt[num_genes] = n[1]
        num_genes = num_genes - len(n[0].child_nodes()) + 1
    # num_alleles_list = sorted(num_genes_wt.keys(), reverse=True)
    return num_genes_wt

def log_probability_of_coalescent_frames(coalescent_frames, haploid_pop_size):
    """
    Under the classical neutral coalescent \citep{Kingman1982,
    Kingman1982b}, the waiting times between coalescent events in a
    sample of $k$ alleles segregating in a  population of (haploid) size
    $N_e$ is distributed exponentially with a rate parameter of
    :math`\\frac{{k \choose 2}}{N_e}`::

        .. math::

            \\Pr(T) =  \\frac{{k \\choose 2}}{N_e} \\e{-  \\frac{{k \\choose 2}}{N_e} T},

    where $T$ is the length of  (chronological) time in which there are
    $k$ alleles in the sample (i.e., for $k$ alleles to coalesce into
    $k-1$ alleles).
    """
    lp = 0.0
    for k, t in coalescent_frames.items():
        k2N = (float(k * (k-1)) / 2) / haploid_pop_size
#         k2N = float(combinatorics.choose(k, 2)) / haploid_pop_size
        lp =  lp + math.log(k2N) - (k2N * t)
    return lp

def log_probability_of_coalescent_tree(tree, haploid_pop_size, ultrametricity_precision=constants.DEFAULT_ULTRAMETRICITY_PRECISION):
    """
    Wraps up extraction of coalescent frames and reporting of probability.
    """
    return log_probability_of_coalescent_frames(extract_coalescent_frames(tree),
            haploid_pop_size)

###############################################################################
## Tree Simulations

def contained_coalescent_tree(containing_tree,
        gene_to_containing_taxon_map,
        edge_pop_size_attr="pop_size",
        default_pop_size=1,
        rng=None):
    """
    Returns a gene tree simulated under the coalescent contained within a
    population or species tree.

        ``containing_tree``
            The population or species tree. If ``edge_pop_size_map`` is not None,
            and population sizes given are non-trivial (i.e., >1), then edge
            lengths on this tree are in units of generations. Otherwise edge
            lengths are in population units; i.e. 2N generations for diploid
            populations of size N, or N generations for diploid populations of
            size N.

        ``gene_to_containing_taxon_map``
            A TaxonNamespaceMapping object mapping Taxon objects in the
            ``containing_tree`` TaxonNamespace to corresponding Taxon objects in the
            resulting gene tree.

        ``edge_pop_size_attr``
            Name of attribute of edges that specify population size. By default
            this is "pop_size". If this attribute does not exist,
            ``default_pop_size`` will be used.  The value for this attribute
            should be the haploid population size or the number of genes;
            i.e.  2N for a diploid population of N individuals, or N for a
            haploid population of N individuals. This value determines how
            branch length units are interpreted in the input tree,
            ``containing_tree``.  If a biologically-meaningful value, then branch
            lengths on the ``containing_tree`` are properly read as generations.
            If not (e.g. 1 or 0), then they are in population units, i.e. where
            1 unit of time equals 2N generations for a diploid population of
            size N, or N generations for a haploid population of size N.
            Otherwise time is in generations. If this argument is None, then
            population sizes default to ``default_pop_size``.

        ``default_pop_size``
            Population size to use if ``edge_pop_size_attr`` is None or
            if an edge does not have the attribute. Defaults to 1.

    The returned gene tree will have the following extra attributes:

        ``pop_node_genes``
            A dictionary with nodes of ``containing_tree`` as keys and a list of gene
            tree nodes that are uncoalesced as values.

    Note that this function does very much the same thing as
    ``constrained_kingman()``, but provides a very different API.
    """

    if rng is None:
        rng = GLOBAL_RNG

    gene_tree_taxon_namespace = gene_to_containing_taxon_map.domain_taxon_namespace
    if gene_tree_taxon_namespace is None:
        gene_tree_taxon_namespace = dendropy.TaxonNamespace()
        for gene_taxa in pop_gene_taxa_map:
            for taxon in gene_taxa:
                gene_tree_taxon_namespace.add(taxon)
    gene_tree = dendropy.Tree(taxon_namespace=gene_tree_taxon_namespace)
    gene_tree.is_rooted = True

    pop_node_genes = {}
    pop_gene_taxa = gene_to_containing_taxon_map.reverse
    for nd in containing_tree.postorder_node_iter():
        if nd.taxon and nd.taxon in pop_gene_taxa:
            pop_node_genes[nd] = []
            gene_taxa = pop_gene_taxa[nd.taxon]
            for gene_taxon in gene_taxa:
                gene_node = dendropy.Node()
                gene_node.taxon = gene_taxon
                pop_node_genes[nd].append(gene_node)
            #gene_nodes = [dendropy.Node() for i in range(len(gene_taxa))]
            #for gidx, gene_node in enumerate(gene_nodes):
            #    gene_node.taxon = gene_taxa[gidx]
            #    pop_node_genes[nd].append(gene_node)

    for edge in containing_tree.postorder_edge_iter():

        if edge_pop_size_attr and hasattr(edge, edge_pop_size_attr):
            pop_size = getattr(edge, edge_pop_size_attr)
        else:
            pop_size = default_pop_size
        if edge.head_node.parent_node is None:
            if len(pop_node_genes[edge.head_node]) > 1:
                final = coalesce_nodes(nodes=pop_node_genes[edge.head_node],
                                            pop_size=default_pop_size,
                                            period=None,
                                            rng=rng)
            else:
                final = pop_node_genes[edge.head_node]
            gene_tree.seed_node = final[0]
        else:
            uncoal = coalesce_nodes(nodes=pop_node_genes[edge.head_node],
                                         pop_size=pop_size,
                                         period=edge.length,
                                         rng=rng)
            if edge.tail_node not in pop_node_genes:
                pop_node_genes[edge.tail_node] = []
            pop_node_genes[edge.tail_node].extend(uncoal)

    gene_tree.pop_node_genes = pop_node_genes
    return gene_tree

def pure_kingman_tree(taxon_namespace, pop_size=1, rng=None):
    """
    Generates a tree under the unconstrained Kingman's coalescent process.

    Parameters
    ----------
    taxon_namespace: |TaxonNamespace| instance
        A pre-populated |TaxonNamespace| where the contained |Taxon| instances
        represent the genes or individuals sampled from the population.
    pop_size : numeric
        The size of the population from the which the coalescent process is
        sampled.


    Returns
    -------
    t : |Tree|
        A tree sampled from the Kingman's neutral coalescent.

    """
    if rng is None:
        rng = GLOBAL_RNG # use the global rng by default
    nodes = [dendropy.Node(taxon=t) for t in taxon_namespace]
    seed_node = coalesce_nodes(nodes=nodes,
                                    pop_size=pop_size,
                                    period=None,
                                    rng=rng,
                                    use_expected_tmrca=False)[0]
    tree = dendropy.Tree(taxon_namespace=taxon_namespace, seed_node=seed_node)
    return tree

def pure_kingman_tree_shape(num_leaves, pop_size=1, rng=None):
    """
    Like :func:`dendropy.model.pure_kingman_tree`, but does not assign taxa to tips.

    Parameters
    ----------
    num_leaves : int
        Number of individuals/genes sampled.
    pop_size : numeric
        The size of the population from the which the coalescent process is
        sampled.

    Returns
    -------
    t : |Tree|
        A tree sampled from the Kingman's neutral coalescent.

    """
    if rng is None:
        rng = GLOBAL_RNG # use the global rng by default
    nodes = [dendropy.Node() for t in range(num_leaves)]
    seed_node = coalesce_nodes(nodes=nodes,
                                    pop_size=pop_size,
                                    period=None,
                                    rng=rng,
                                    use_expected_tmrca=False)[0]
    tree = dendropy.Tree(seed_node=seed_node)
    return tree

def mean_kingman_tree(taxon_namespace, pop_size=1, rng=None):
    """
    Returns a tree with coalescent intervals given by the expected times under
    Kingman's neutral coalescent.
    """
    if rng is None:
        rng = GLOBAL_RNG # use the global rng by default
    nodes = [dendropy.Node(taxon=t) for t in taxon_namespace]
    seed_node = coalesce_nodes(nodes=nodes,
                                    pop_size=pop_size,
                                    period=None,
                                    rng=rng,
                                    use_expected_tmrca=True)[0]
    tree = dendropy.Tree(taxon_namespace=taxon_namespace, seed_node=seed_node)
    return tree

def constrained_kingman_tree(pop_tree,
                        gene_tree_list=None,
                        rng=None,
                        gene_node_label_fn=None,
                        num_genes_attr='num_genes',
                        pop_size_attr='pop_size',
                        decorate_original_tree=False):
    """
    Given a population tree, ``pop_tree`` this will return a *pair of
    trees*: a gene tree simulated on this population tree based on
    Kingman's n-coalescent, and population tree with the additional
    attribute 'gene_nodes' on each node, which is a list of
    uncoalesced nodes from the gene tree associated with the given
    node from the population tree.

    ``pop_tree`` should be a DendroPy Tree object or an object
    of a class derived from this with the following attribute
    ``num_genes`` -- the number of gene samples from each population in the
    present.  Each edge on the tree should also have the attribute

    ``pop_size_attr`` is the attribute name of the edges of ``pop_tree`` that
    specify the population size. By default it is ``pop_size``. The should
    specify the effective *haploid* population size; i.e., number of gene
    in the population: 2 * N in a diploid population of N individuals,
    or N in a haploid population of N individuals.

    If ``pop_size`` is 1 or 0 or None, then the edge lengths of ``pop_tree`` is
    taken to be in haploid population units; i.e. where 1 unit equals 2N
    generations for a diploid population of size N, or N generations for a
    haploid population of size N. Otherwise the edge lengths of ``pop_tree`` is
    taken to be in generations.

    If ``gene_tree_list`` is given, then the gene tree is added to the
    tree block, and the tree block's taxa block will be used to manage
    the gene tree's ``taxa``.

    ``gene_node_label_fn`` is a function that takes two arguments (a string
    and an integer, respectively, where the string is the containing species
    taxon label and the integer is the gene index) and returns a label for
    the corresponding the gene node.

    if ``decorate_original_tree`` is True, then the list of uncoalesced nodes at
    each node of the population tree is added to the original (input) population
    tree instead of a copy.

    Note that this function does very much the same thing as ``contained_coalescent()``,
    but provides a very different API.
    """

    # get our random number generator
    if rng is None:
        rng = GLOBAL_RNG # use the global rng by default

    if gene_tree_list is not None:
        gtaxa = gene_tree_list.taxon_namespace
    else:
        gtaxa = dendropy.TaxonNamespace()

    if gene_node_label_fn is None:
        gene_node_label_fn = lambda x, y: "%s_%02d" % (x, y)

    # we create a set of gene nodes for each leaf node on the population
    # tree, and associate those gene nodes to the leaf by assignment
    # of 'taxon'.
    for leaf_count, leaf in enumerate(pop_tree.leaf_node_iter()):
        gene_nodes = []
        for gene_count in range(getattr(leaf, num_genes_attr)):
            gene_node = dendropy.Node()
            gene_node.taxon = gtaxa.require_taxon(label=gene_node_label_fn(leaf.taxon.label, gene_count+1))
            gene_nodes.append(gene_node)
        leaf.gene_nodes = gene_nodes

    # We iterate through the edges of the population tree in post-order,
    # i.e., visiting child edges before we visit parent edges. For
    # each edge visited, we take the genes found in the child nodes,
    # and run the coalescent simulation on them attacheded by the length
    # of the edge. Any genes that have not yet coalesced at the end of
    # this period are added to the genes of the tail (parent) node of
    # the edge.

    if decorate_original_tree:
        working_poptree = pop_tree
    else:
        # start with a new (deep) copy of the population tree so as to not
        # to change the original tree
        working_poptree = dendropy.Tree(pop_tree)

    # start with a new tree
    gene_tree = dendropy.Tree()
    gene_tree.taxon_namespace = gtaxa
    for edge in working_poptree.postorder_edge_iter():

        # if mrca root, run unconstrained coalescent
        if edge.head_node.parent_node is None:
            if len(edge.head_node.gene_nodes) > 1:
                final = coalesce_nodes(nodes=edge.head_node.gene_nodes,
                                            pop_size=pop_size,
                                            period=None,
                                            rng=rng)
            else:
                final = edge.head_node.gene_nodes
            gene_tree.seed_node = final[0]
        else:

            if hasattr(edge, pop_size_attr):
                pop_size = getattr(edge, pop_size_attr)
            else:
                # this means all our time will be in population units
                pop_size = 1

            uncoal = coalesce_nodes(nodes=edge.head_node.gene_nodes,
                                         pop_size=pop_size,
                                         period=edge.length,
                                         rng=rng)
            if not hasattr(edge.tail_node, 'gene_nodes'):
                edge.tail_node.gene_nodes = []
            edge.tail_node.gene_nodes.extend(uncoal)

    gene_tree.is_rooted = True
    if gene_tree_list is not None:
        gene_tree_list.append(gene_tree)
        return gene_tree, working_poptree
    else:
        return gene_tree, working_poptree