#! /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.
##
##############################################################################

"""
Models, modeling and model-fitting of the protracted speciation, as described in::

        Etienne, R.S., Morlon, H., and Lambert, A. 2014. Estimating the
        duration of speciation from phylogenies. Evolution 2014: 2430-2440.
        doi:10.1111/evo.12433

"""

import math
import itertools
import dendropy
from dendropy.utility import GLOBAL_RNG
from dendropy.utility.error import ProcessFailedException
from dendropy.utility.error import TreeSimTotalExtinctionException
from dendropy.calculate import probability


def _D(speciation_initiation_rate,
       speciation_completion_rate,
       incipient_species_extinction_rate):
    """
    Returns value of D, as given in eq. 5 in Etienne et al.
    (2014).

    Parameters
    ----------

    speciation_initiation_rate : float
        The birth rate, b (the incipient species birth
        rate and the "good" species birth rate are assumed to be equal):
        the rate at which new (incipient) species are produced from
        either incipient or "good" species lineages.
    speciation_completion_rate : float
        The rate at which incipient species get converted to good
        species, $\lambda_1$.
    incipient_species_extinction_rate : float
        The incipient species exctinction rate, $\mu_1$: the rate at which
        incipient species go extinct.

    Returns
    -------
    t : float
        The duration of speciation.

    """
    D = math.sqrt(
            pow(speciation_completion_rate + speciation_initiation_rate - incipient_species_extinction_rate, 2)
            + (4.0 * speciation_completion_rate * incipient_species_extinction_rate)
        )
    return D

def _phi(speciation_initiation_rate,
       speciation_completion_rate,
       incipient_species_extinction_rate):
    """
    Returns value of $\varphi$, as given in eq. 6 in Etienne et al.
    (2014).

    Parameters
    ----------

    speciation_initiation_rate : float
        The birth rate, b (the incipient species birth
        rate and the "good" species birth rate are assumed to be equal):
        the rate at which new (incipient) species are produced from
        either incipient or "good" species lineages.
    speciation_completion_rate : float
        The rate at which incipient species get converted to good
        species, $\lambda_1$.
    incipient_species_extinction_rate : float
        The incipient species exctinction rate, $\mu_1$: the rate at which
        incipient species go extinct.

    Returns
    -------
    t : float
        The duration of speciation.

    """
    phi = speciation_completion_rate - speciation_initiation_rate + incipient_species_extinction_rate
    return phi

def expected_duration_of_speciation(
        speciation_initiation_rate,
        speciation_completion_rate,
        incipient_species_extinction_rate,
        D=None,
        ):
    """
    Returns mean duration of speciation, following Eqs. 4 in Etienne et al.
    (2014):

        The duration of speciation differs from the speciation-completion
        time in that the latter is the waiting time until a single
        incipient lineage completes the speciation process if extinction
        was zero, whereas the former is the time needed for an incipient
        species or one of its descendants to complete speciation, condi-
        tional on the fact that speciation completes, that is, this is the
        time taken by any species that succeeded in speciating completely.

    Parameters
    ----------

    speciation_initiation_rate : float
        The birth rate, b (the incipient species birth
        rate and the "good" species birth rate are assumed to be equal):
        the rate at which new (incipient) species are produced from
        either incipient or "good" species lineages.
    speciation_completion_rate : float
        The rate at which incipient species get converted to good
        species, $\lambda_1$.
    incipient_species_extinction_rate : float
        The incipient species exctinction rate, $\mu_1$: the rate at which
        incipient species go extinct.
    D : float
        Value of ``D`` (as given in Eq. 5 in Etienne et al. 2014). Will be
        calculated if not specified.

    Returns
    -------
    t : float
        The duration of speciation.

    """
    if D is None:
        D = _D(
            speciation_initiation_rate=speciation_initiation_rate,
            speciation_completion_rate=speciation_completion_rate,
            incipient_species_extinction_rate=incipient_species_extinction_rate)
    t1 = 2.0/(D - speciation_completion_rate + speciation_initiation_rate - incipient_species_extinction_rate )
    t2 = math.log(2.0/(1+((speciation_completion_rate - speciation_initiation_rate + incipient_species_extinction_rate)/D)))
    t = t1 * t2
    return t

def probability_of_duration_of_speciation(
        tau,
        speciation_initiation_rate,
        speciation_completion_rate,
        incipient_species_extinction_rate,
        D=None,
        phi=None,
        ):
    """
    Returns probability of duration of speciation, tau, following Eqs. 6
    in Etienne et al.

    Parameters
    ----------

    tau : float
        The duration of speciation.
    speciation_initiation_rate : float
        The birth rate, b (the incipient species birth
        rate and the "good" species birth rate are assumed to be equal):
        the rate at which new (incipient) species are produced from
        either incipient or "good" species lineages.
    speciation_completion_rate : float
        The rate at which incipient species get converted to good
        species, $\lambda_1$.
    incipient_species_extinction_rate : float
        The incipient species exctinction rate, $\mu_1$: the rate at which
        incipient species go extinct.
    D : float
        Value of ``D`` (as given in Eq. 5 in Etienne et al. 2014). Will be
        calculated if not specified.
    phi : float
        Value of ``phi`` (as given in Eq. 7 in Etienne et al. 2014). Will be
        calculated if not specified.

    Returns
    -------
    p : float
        The probability of the duration of speciation, tau.

    """
    if D is None:
        D = _D(
            speciation_initiation_rate=speciation_initiation_rate,
            speciation_completion_rate=speciation_completion_rate,
            incipient_species_extinction_rate=incipient_species_extinction_rate)
    if phi is None:
        phi = _phi(
            speciation_initiation_rate=speciation_initiation_rate,
            speciation_completion_rate=speciation_completion_rate,
            incipient_species_extinction_rate=incipient_species_extinction_rate)
    n1 = 2.0 * pow(D, 2) * math.exp(-D * tau) * (D + phi)
    d1 = pow(D + phi + math.exp(-D * tau) * (D-phi), 2)
    return n1/d1

def log_probability_of_duration_of_speciation(
        tau,
        speciation_initiation_rate,
        speciation_completion_rate,
        incipient_species_extinction_rate,
        D=None,
        phi=None,
        ):
    """
    Returns probability of duration of speciation, tau, following Eqs. 6
    in Etienne et al.

    Parameters
    ----------

    tau : float
        The duration of speciation.
    speciation_initiation_rate : float
        The birth rate, b (the incipient species birth
        rate and the "good" species birth rate are assumed to be equal):
        the rate at which new (incipient) species are produced from
        either incipient or "good" species lineages.
    speciation_completion_rate : float
        The rate at which incipient species get converted to good
        species, $\lambda_1$.
    incipient_species_extinction_rate : float
        The incipient species exctinction rate, $\mu_1$: the rate at which
        incipient species go extinct.
    D : float
        Value of ``D`` (as given in Eq. 5 in Etienne et al. 2014). Will be
        calculated if not specified.
    phi : float
        Value of ``phi`` (as given in Eq. 7 in Etienne et al. 2014). Will be
        calculated if not specified.

    Returns
    -------
    p : float
        The probability of the duration of speciation, tau.

    """
    if D is None:
        D = _D(
            speciation_initiation_rate=speciation_initiation_rate,
            speciation_completion_rate=speciation_completion_rate,
            incipient_species_extinction_rate=incipient_species_extinction_rate)
    if phi is None:
        phi = _phi(
            speciation_initiation_rate=speciation_initiation_rate,
            speciation_completion_rate=speciation_completion_rate,
            incipient_species_extinction_rate=incipient_species_extinction_rate)
    n1 = math.log(2.0) + (2 * math.log(D)) - (D * tau) + math.log(D + phi)
    d1 = 2 * (math.log(D + phi + math.exp(-D * tau)*(D-phi)))
    return n1 - d1

def maximum_probability_duration_of_speciation(
        speciation_initiation_rate,
        speciation_completion_rate,
        incipient_species_extinction_rate,
        D=None,
        phi=None,
        ):
    """
    Returns duration of speciation that maximizes probability under given
    process parameters, following eq. 8 of Etienne et al (2014).

    Parameters
    ----------

    speciation_initiation_rate : float
        The birth rate, b (the incipient species birth
        rate and the "good" species birth rate are assumed to be equal):
        the rate at which new (incipient) species are produced from
        either incipient or "good" species lineages.
    speciation_completion_rate : float
        The rate at which incipient species get converted to good
        species, $\lambda_1$.
    incipient_species_extinction_rate : float
        The incipient species exctinction rate, $\mu_1$: the rate at which
        incipient species go extinct.
    D : float
        Value of ``D`` (as given in Eq. 5 in Etienne et al. 2014). Will be
        calculated if not specified.
    phi : float
        Value of ``phi`` (as given in Eq. 7 in Etienne et al. 2014). Will be
        calculated if not specified.

    Returns
    -------
    t : float
        The duration of speciation with the maximum probability under the
        given process parameters.

    """
    if D is None:
        D = _D(
            speciation_initiation_rate=speciation_initiation_rate,
            speciation_completion_rate=speciation_completion_rate,
            incipient_species_extinction_rate=incipient_species_extinction_rate)
    if phi is None:
        phi = _phi(
            speciation_initiation_rate=speciation_initiation_rate,
            speciation_completion_rate=speciation_completion_rate,
            incipient_species_extinction_rate=incipient_species_extinction_rate)
    x = 1.0/D * math.log((D-phi)/(D+phi))
    return max(0, x)

class ProtractedSpeciationProcess(object):

    class ProtractedSpeciationProcessLineage(object):

        def __init__(self,
                index,
                parent_lineage,
                speciation_initiation_time,
                is_orthospecies,
                orthospecies_index):
            self._index = index
            self.is_orthospecies = is_orthospecies
            self.parent_lineage = parent_lineage
            self.speciation_initiation_time = speciation_initiation_time
            self.speciation_completion_time = None
            self.extinction_time = None
            self.lineage_tree_node_history = []
            self._label = "L{}".format(self._index)
            self.orthospecies_index = orthospecies_index

        def _get_node(self):
            return self.lineage_tree_node_history[-1]
        def _set_node(self, node):
            self.lineage_tree_node_history.append(node)
        node = property(_get_node, _set_node)

        def _get_index(self):
            return self._index
        index = property(_get_index)

        def _get_label(self):
            return self._label
        label = property(_get_label)

        def __str__(self):
            return self._get_label()

    def __init__(self,
            speciation_initiation_from_orthospecies_rate,
            speciation_initiation_from_incipient_species_rate,
            speciation_completion_rate,
            orthospecies_extinction_rate,
            incipient_species_extinction_rate,
            rng=None,
            ):
        self.speciation_initiation_from_orthospecies_rate = speciation_initiation_from_orthospecies_rate
        self.orthospecies_extinction_rate = orthospecies_extinction_rate
        self.speciation_initiation_from_incipient_species_rate = speciation_initiation_from_incipient_species_rate
        self.speciation_completion_rate = speciation_completion_rate
        self.incipient_species_extinction_rate = incipient_species_extinction_rate
        if rng is None:
            self.rng = GLOBAL_RNG
        else:
            self.rng = rng
        self.tree_factory = dendropy.Tree
        self.node_factory = dendropy.Node
        self.reset()

    def reset(self):
        self.current_time = 0.0
        self.current_lineage_index = 0
        self.current_orthospecies_index = 0
        self.current_node_index = 0
        self.current_orthospecies_lineages = []
        self.current_incipient_species_lineages = []
        self.lineage_to_orthospecies_tree_node_map = {}

    def generate_sample(self, **kwargs):
        """

        Samples from the Protracted Speciation Model process, returning a tuple:

            -   the lineage tree: this tree has all nodes/lineages, i.e. both
                "good" species as well as incipient species.
            -   the (ortho- or confirmed- or "good"-)species tree: the tree
                only has "good" species, i.e. with all incipient species
                pruned out.

        Each node on the protracted speciation tree as will as the "good" species
        tree will have an attribute, ``protracted_speciation_model_lineage``,
        which is a reference to a
        :class:`~dendropy.model.birthdeath.ProtractedSpeciationProcess.ProtractedSpeciationProcessLineage`
        instance which represents the lineage associated with this node. Note
        that each node can only be associated with a single lineage, but a
        lineage might span several nodes.

        If ``is_correlate_lineage_and_species_trees`` is |True|,
        then additional attributes will be added. See
        :meth:``dendropy.model.protractedspeciation.ProtractedSpeciationProcess.correlate_lineage_and_species_trees`
        for details.

        Parameters
        ----------

        max_time : float or |None|
            Terminate and return results when this time is reached. If |None|,
            then do not terminated based on run time.
        max_extant_orthospecies : int or |None|
            Terminate and return results when this number of tips are found in
            the confirmed-species tree (i.e., the pruned tree consisting of only
            "good" species). If |None|, then do not terminate
            based on the number of tipes on the confirmed-species tree.
        max_extant_lineages : int or |None|
            Terminate and return results when this number of tips are found in
            the lineage tree (i.e. the tree with both incipient and good
            species). If |None|, then do not terminate based on the
            number of tipes on the incipient species tree.
        is_initial_lineage_orthospecies : bool
            Whether the first lineage that initialies the process is a
            "good" species or not. Defaults to |True|: first species on
            the tree is a "good" species.
        is_retry_on_total_extinction : bool
            If |False|, then a TreeSimTotalExtinctionException will be raised
            if all lineages go extinct before the termination conditions are
            met. Defaults to |True|: if all lineages go extinct before the
            termination conditions are met, then the simulation is rerun, up to
            a maximum of ``max_retries``.
        max_retries : int
            Maximum number of runs to execute in the event of
            prematurely-terminated simulations due to all lineages going
            extinct. Once this number or re-runs is exceed, then
            TreeSimTotalExtinctionException is raised. Defaults to 1000. Set to
            |None| to never quit trying.
        is_correlate_lineage_and_species_trees: bool
            [NOT USED]
            If |True| then additional attributes will be added to the
            resulting trees to relate them. See
            :meth:``dendropy.model.protractedspeciation.ProtractedSpeciationProcess.correlate_lineage_and_species_trees`
            for details.

        Returns
        -------
        lineage_tree : |Tree| instance
            A tree from the protracted speciation process, with all lineages
            (good species as well as incipient species).
        orthospecies_tree : |Tree| instance
            A tree from the protracted speciation process with only "good" species.

        """
        is_retry_on_total_extinction = kwargs.pop("is_retry_on_total_extinction", True)
        max_retries = kwargs.pop("max_retries", 1000)
        num_retries = 0
        lineage_tree = None
        orthospecies_tree = None
        while True:
            try:
                lineage_tree, orthospecies_tree = self._run_protracted_speciation_process(**kwargs)
                break
            except ProcessFailedException:
                if not is_retry_on_total_extinction:
                    raise
                num_retries += 1
                if max_retries is not None and num_retries > max_retries:
                    raise
        assert lineage_tree is not None
        return lineage_tree, orthospecies_tree

    # def correlate_lineage_and_species_trees(self,
    #         lineage_tree,
    #         orthospecies_tree):
    #     """
    #     Correlates the protracted speciation tree and the corresponding pruned
    #     "good" species tree from a single sample of the protracted speciation
    #     process (i.e., a call to ``generate_sample()``).

    #     Each node on the lineage tree will have the following
    #     attributes added:

    #         - ``is_orthospeciation_event`` : |True| if the node represents a
    #          "good" speciation event, |False| otherwise.

    #     Each internal node on the "good" species tree will have the following
    #     attributes added:

    #         - ``lineage_tree_node``: a reference to the node on
    #           the protratcted speciation tree to which it corresponds.

    #     Each leaf node on the orthospecies tree will have the following
    #     attributes added:

    #         - ``included_lineage_tree_leaf_nodes``: the set of
    #           terminal/leaf nodes on the protracted speciation tree which are
    #           descended/included in it.

    #     """
    #     return lineage_tree, orthospecies_tree
    #     lineage_tree.calc_node_ages()

    #     orthospecies_tree_leaf_node_parent_age_to_desc_node_and_lineages_map = {}
    #     for os_nd in orthospecies_tree.leaf_node_iter():
    #         if os_nd.parent_node is None:
    #             raise ProcessFailedException()
    #         age = str(os_nd.parent_node.age)
    #         orthospecies_tree_leaf_node_parent_age_to_desc_node_and_lineages_map[age] = []
    #         for ch1 in os_nd.parent_node.child_node_iter():
    #             if not ch1._child_nodes:
    #                 # desc_lineages = set([sch1.protracted_speciation_model_lineage for sch1 in ch1.preorder_node_iter])
    #                 orthospecies_tree_leaf_node_parent_age_to_desc_node_and_lineages_map[age].append((ch1, ch1.protracted_speciation_model_lineage))


    #             # desc_lineages = set([sch1.protracted_speciation_model_lineage for sch1 in ch1.preorder_node_iter])
    #             # orthospecies_tree_leaf_node_parent_age_to_desc_node_and_lineages_map.append( (ch1, desc_lineages) )

    #     for ln_nd in lineage_tree.leaf_node_iter():
    #         cur_nd = ln_nd
    #         if cur_nd.parent_node is None:
    #             raise ProcessFailedException()
    #         focal_nd_siblings = [snd for snd in ln_nd.parent_node.child_node_iter() if snd is not ln_nd]
    #         focal_nd_parents = [pnd for pnd in ln_nd.ancestor_iter(inclusive=False)]
    #         chain = []
    #         age_found = False
    #         found = None
    #         while True:
    #             if cur_nd is None:
    #                 print("!!! {}: current is None".format(" > ".join(chain)) )
    #                 break
    #             chain.append(cur_nd.label)
    #             if cur_nd.parent_node is None:
    #                 print("!!! {}: parent is None".format(" > ".join(chain)) )
    #                 break
    #             age = str(cur_nd.parent_node.age)
    #             if age in orthospecies_tree_leaf_node_parent_age_to_desc_node_and_lineages_map:
    #                 age_found = True
    #                 cur_nd_desc_lineages = set([sch.protracted_speciation_model_lineage for sch in cur_nd.preorder_iter()])
    #                 for orthospecies_nd, orthospecies_lineage in orthospecies_tree_leaf_node_parent_age_to_desc_node_and_lineages_map[age]:
    #                     if orthospecies_lineage in cur_nd_desc_lineages:
    #                         print(":::: {} found in desc of {}: {}".format(orthospecies_nd.protracted_speciation_model_lineage.label, ln_nd.label, [x.label for x in cur_nd_desc_lineages]))
    #                         orthospecies_nd.included_lineage_tree_leaf_nodes.add(ln_nd)
    #                         found = orthospecies_nd
    #                         break
    #                 break
    #             cur_nd = cur_nd.parent_node
    #         if not age_found:
    #             print("Age not found: {}".format(ln_nd.label))
    #         elif not found:
    #             print("Not found: {}".format(ln_nd.label))
    #         else:
    #             print("Found: {} => {}".format(ln_nd.label, found.label))
    #     # lineage_tree_node_to_orthospecies_leaf_map = {}
    #     # lineage_orthospecies_leaf_map = {}
    #     # for orthospecies_tree_nd in orthospecies_tree:
    #     #     if not orthospecies_tree_nd._child_nodes:
    #     #         lineage_orthospecies_leaf_map[orthospecies_tree_nd.protracted_speciation_model_lineage] = orthospecies_tree_nd
    #     #         continue
    #     #     lineage_tree_node = sorted(orthospecies_tree_nd.protracted_speciation_model_lineage.lineage_tree_node_history, key=lambda x: x.age)[-1]
    #     #     for history_nd in orthospecies_tree_nd.protracted_speciation_model_lineage.lineage_tree_node_history:
    #     #         if not orthospecies_tree_nd._child_nodes:
    #     #             lineage_tree_node_to_orthospecies_leaf_map[history_nd] = orthospecies_tree_nd
    #     #         # if history_nd.age is not None and history_nd.age >= lineage_tree_node.age:
    #     #         #     lineage_tree_node = history_nd
    #     #     lineage_tree_node.is_orthospeciation_event = True
    #     #     orthospecies_tree_nd.lineage_tree_node = lineage_tree_node

    #     # for lineage_tree_leaf in lineage_tree.leaf_node_iter():
    #     #     nd = lineage_tree_leaf
    #     #     print("#### {}: {}".format(lineage_tree_leaf.label, nd.protracted_speciation_model_lineage.is_orthospecies))
    #     #     while nd and nd.protracted_speciation_model_lineage.speciation_completion_time is not None:
    #     #         print(">>>> {}: {}".format(lineage_tree_leaf.label, nd.label))
    #     #         nd = nd.parent_node
    #     #     if not nd:
    #     #         print("!!! Failed: {}".format(lineage_tree_leaf.label))
    #     #     elif not nd.protracted_speciation_model_lineage.is_orthospecies:
    #     #         print("!!! Failed 2: {}".format(lineage_tree_leaf.label))
    #     #     else:
    #     #         current_lineages = [desc.protracted_speciation_model_lineage for desc in nd.preorder_iter()]
    #     #         print("set for {}: {}".format(lineage_tree_leaf.label, [x.label for x in current_lineages]))
    #     #         for lineage in current_lineages:
    #     #             if lineage in lineage_orthospecies_leaf_map:
    #     #                 lineage_orthospecies_leaf_map[lineage].included_lineage_tree_leaf_nodes.append(lineage_tree_leaf)
    #     #                 break
    #     #         else:
    #     #             print("Not found: {}".format(lineage_tree_leaf.label))


    #     # for lineage_tree_leaf in lineage_tree.leaf_node_iter():
    #     #     lineage = lineage_tree_leaf.protracted_speciation_model_lineage
    #     #     while lineage is not None and not lineage.is_orthospecies:
    #     #         lineage = lineage.parent_lineage
    #     #     if lineage is None:
    #     #         ## TODO: special case
    #     #         pass
    #     #     for orthospecies_tree_nd in orthospecies_tree:
    #     #         if orthospecies_tree_nd.protracted_speciation_model_lineage is lineage:
    #     #             if orthospecies_tree_nd._child_nodes:
    #     #                 print("Internal node {} for: {}".format(orthospecies_tree_nd.label, lineage_tree_leaf.label))
    #     #             else:
    #     #                 print("OK {} for: {}".format(orthospecies_tree_nd.label, lineage_tree_leaf.label))
    #     #             break
    #     #     else:
    #     #         print("Not found: {}".format(lineage_tree_leaf.label))



    #     # for lineage_tree_leaf in lineage_tree.leaf_node_iter():
    #     #     print("....{}".format(lineage_tree_leaf.label))
    #     #     if lineage_tree_leaf in lineage_tree_node_to_orthospecies_leaf_map:
    #     #         print("516: {}: {} => {}".format(lineage_tree_node_to_orthospecies_leaf_map[lineage_tree_leaf].label, lineage_tree_leaf.label, list(x.label for x in lineage_tree_node_to_orthospecies_leaf_map[lineage_tree_leaf].included_lineage_tree_leaf_nodes)))
    #     #         lineage_tree_node_to_orthospecies_leaf_map[lineage_tree_leaf].included_lineage_tree_leaf_nodes.append(lineage_tree_leaf)
    #     #         print("518: {}: {} => {}".format(lineage_tree_node_to_orthospecies_leaf_map[lineage_tree_leaf].label, lineage_tree_leaf.label, list(x.label for x in lineage_tree_node_to_orthospecies_leaf_map[lineage_tree_leaf].included_lineage_tree_leaf_nodes)))
    #     #     else:
    #     #         print("519")
    #     #         lineage = lineage_tree_leaf.protracted_speciation_model_lineage
    #     #         while lineage is not None and lineage not in lineage_orthospecies_leaf_map:
    #     #             lineage = lineage.parent_lineage
    #     #         if lineage is None:
    #     #             print("524")
    #     #             orthospecies_tree.seed_node.included_lineage_tree_leaf_nodes.append(lineage_tree_leaf)
    #     #             print("<<< {} >>>".format(lineage_tree_leaf.label))
    #     #         else:
    #     #             print("528")
    #     #             lineage_orthospecies_leaf_map[lineage].included_lineage_tree_leaf_nodes.append(lineage_tree_leaf)

    #         # # while lineage is not None and not lineage.is_orthospecies:
    #         # #     lineage = lineage.parent_lineage
    #         # if lineage is None:
    #         #     orthospecies_tree_nd = orthospecies_tree.seed_node
    #         # else:
    #         #     orthospecies_tree_nd = self.lineage_to_orthospecies_tree_node_map[lineage]
    #         #     # for xnd in orthospecies_tree.leaf_node_iter():
    #         #     #     if xnd.protracted_speciation_model_lineage is lineage:
    #         #     #         orthospecies_tree_nd = xnd
    #         #     #         break
    #         #     # else:
    #         #     #     orthospecies_tree_nd = orthospecies_tree.seed_node
    #         # try:
    #         #     orthospecies_tree_nd.included_lineage_tree_leaf_nodes.append(lineage)
    #         # except AttributeError:
    #         #     orthospecies_tree_nd.included_lineage_tree_leaf_nodes = [lineage]

    #         # while lineage is not None and lineage not in self.lineage_to_orthospecies_tree_node_map:
    #         #     lineage = lineage.parent_lineage
    #         # if lineage not in self.lineage_to_orthospecies_tree_node_map:
    #         #     orthospecies_tree_nd = orthospecies_tree.seed_node
    #         # else:
    #         #     orthospecies_tree_nd = self.lineage_to_orthospecies_tree_node_map[lineage]
    #         # orthospecies_tree_nd.included_lineage_tree_leaf_nodes.append(lineage)

    #     return lineage_tree, orthospecies_tree

    def _run_protracted_speciation_process(self, **kwargs):
        self.reset()
        max_time = kwargs.get("max_time", None)
        max_extant_lineages = kwargs.get("max_extant_lineages", None)
        max_extant_orthospecies = kwargs.get("max_extant_orthospecies", None)
        is_correlate_lineage_and_species_trees = kwargs.get("is_correlate_lineage_and_species_trees", False)
        taxon_namespace = kwargs.get("taxon_namespace", None)
        initial_lineage = self._new_lineage(
                parent_lineage=None,
                orthospecies_index=self.current_orthospecies_index,
                is_orthospecies=kwargs.get("is_initial_lineage_orthospecies", True),
                )
        seed_node = self._new_node(lineage=initial_lineage)
        lineage_tree = self.tree_factory( taxon_namespace=taxon_namespace, seed_node=seed_node)
        lineage_tree.is_rooted = True

        while True:

            ## Draw time to next event
            event_rates = []
            num_orthospecies = len(self.current_orthospecies_lineages)
            if max_extant_orthospecies is not None:
                ## note: expensive operation to count leaves!
                try:
                    orthospecies_tree = self._assemble_orthospecies_tree(taxon_namespace=taxon_namespace)
                    num_leaves = len(orthospecies_tree.leaf_nodes())
                    if num_leaves >= max_extant_orthospecies:
                        return self._postprocess_psm_and_orthospecies_trees(
                                orthospecies_tree=orthospecies_tree,
                                lineage_tree=lineage_tree,
                                is_correlate_lineage_and_species_trees=is_correlate_lineage_and_species_trees,
                                )
                except ProcessFailedException:
                    pass

            num_incipient_species = len(self.current_incipient_species_lineages)
            if max_extant_lineages is not None and (num_incipient_species + num_orthospecies) >= max_extant_lineages:
                break

            # Event type 0
            event_rates.append(self.speciation_initiation_from_orthospecies_rate * num_orthospecies)

            # Event type 1
            event_rates.append(self.orthospecies_extinction_rate * num_orthospecies)

            # Event type 2
            event_rates.append(self.speciation_initiation_from_incipient_species_rate * num_incipient_species)

            # Event type 3
            event_rates.append(self.speciation_completion_rate * num_incipient_species)

            # Event type 4
            event_rates.append(self.incipient_species_extinction_rate * num_incipient_species)

            # All events
            rate_of_any_event = sum(event_rates)

            # Waiting time
            waiting_time = self.rng.expovariate(rate_of_any_event)
            if max_time and (self.current_time + waiting_time) > max_time:
                t = max_time - self.current_time
                for lineage in itertools.chain(self.current_orthospecies_lineages, self.current_incipient_species_lineages):
                    lineage.node.edge.length += t
                self.current_time = max_time
                break
            self.current_time += waiting_time
            for lineage in itertools.chain(self.current_orthospecies_lineages, self.current_incipient_species_lineages):
                lineage.node.edge.length += waiting_time

            # Select event
            event_type_idx = probability.weighted_index_choice(weights=event_rates, rng=self.rng)
            assert (event_type_idx >= 0 and event_type_idx <= 4)
            # print("time {}: {}, selected = {}".format(self.current_time, event_rates, event_type_idx))

            if event_type_idx == 0:
                self._process_initiation_of_speciation_from_orthospecies(lineage_tree)
            elif event_type_idx == 1:
                self._process_orthospecies_extinction(lineage_tree)
            elif event_type_idx == 2:
                self._process_initiation_of_speciation_from_incipient_species(lineage_tree)
            elif event_type_idx == 3:
                self._process_completion_of_specation(lineage_tree)
            elif event_type_idx == 4:
                self._process_incipient_species_extinction(lineage_tree)
            else:
                raise Exception("Unexpected event type index: {}".format(event_type_idx))

            if len(self.current_orthospecies_lineages) + len(self.current_incipient_species_lineages) == 0:
                raise TreeSimTotalExtinctionException()

        orthospecies_tree = self._assemble_orthospecies_tree(taxon_namespace=taxon_namespace)
        return self._postprocess_psm_and_orthospecies_trees(
                lineage_tree=lineage_tree,
                orthospecies_tree=orthospecies_tree,
                is_correlate_lineage_and_species_trees=is_correlate_lineage_and_species_trees,
                )

    def _process_initiation_of_speciation_from_orthospecies(self, tree):
        # parent_lineage = self.rng.choice(self.current_orthospecies_lineages)
        # parent_node = parent_lineage.node
        # new_lineage = self._new_lineage(
        #         parent_lineage=parent_lineage,
        #         orthospecies_index=parent_lineage.orthospecies_index,
        #         is_orthospecies=False)
        # c1 = self._new_node(lineage=parent_lineage)
        # c2 = self._new_node(lineage=new_lineage)
        # parent_node.add_child(c1)
        # parent_node.add_child(c2)
        parent_lineage = self.rng.choice(self.current_orthospecies_lineages)
        self._process_initiation_of_speciation(parent_lineage=parent_lineage)

    def _process_initiation_of_speciation_from_incipient_species(self, tree):
        # parent_lineage = self.rng.choice(self.current_incipient_species_lineages)
        # parent_node = parent_lineage.node
        # new_lineage = self._new_lineage(
        #         parent_lineage=parent_lineage,
        #         orthospecies_index=parent_lineage.orthospecies_index,
        #         is_orthospecies=False)
        # c1 = self._new_node(lineage=parent_lineage)
        # c2 = self._new_node(lineage=new_lineage)
        # parent_node.add_child(c1)
        # parent_node.add_child(c2)
        parent_lineage = self.rng.choice(self.current_incipient_species_lineages)
        self._process_initiation_of_speciation(parent_lineage=parent_lineage)

    def _process_initiation_of_speciation(self, parent_lineage):
        parent_node = parent_lineage.node
        new_lineage = self._new_lineage(
                parent_lineage=parent_lineage,
                orthospecies_index=parent_lineage.orthospecies_index,
                is_orthospecies=False)
        c1 = self._new_node(lineage=parent_lineage)
        c2 = self._new_node(lineage=new_lineage)
        parent_node.add_child(c1)
        parent_node.add_child(c2)

        # parent_node = parent_lineage.node
        # new_lineage1 = self._new_lineage(
        #         parent_lineage=parent_lineage,
        #         orthospecies_index=parent_lineage.orthospecies_index,
        #         is_orthospecies=False,
        #         add_to_current_lineages=True)
        # new_lineage2 = self._new_lineage(
        #         parent_lineage=parent_lineage,
        #         orthospecies_index=parent_lineage.orthospecies_index,
        #         is_orthospecies=parent_lineage.is_orthospecies,
        #         add_to_current_lineages=False)
        # new_lineage2.lineage_tree_node_history = parent_lineage.lineage_tree_node_history
        # c1 = self._new_node(lineage=new_lineage1)
        # c2 = self._new_node(lineage=new_lineage2)
        # parent_node.add_child(c1)
        # parent_node.add_child(c2)

    def _process_completion_of_specation(self, tree):

        lineage = self.rng.choice(self.current_incipient_species_lineages)
        self.current_incipient_species_lineages.remove(lineage)
        self.current_orthospecies_lineages.append(lineage)
        self.current_orthospecies_index += 1
        lineage.orthospecies_index = self.current_orthospecies_index
        lineage.is_orthospecies = True
        lineage.speciation_completion_time = self.current_time

        # original_lineage = self.rng.choice(self.current_incipient_species_lineages)
        # self.current_incipient_species_lineages.remove(original_lineage)
        # self.current_orthospecies_index += 1
        # converted_lineage = self._new_lineage(
        #         parent_lineage=original_lineage,
        #         orthospecies_index=self.current_orthospecies_index,
        #         is_orthospecies=True,
        #         add_to_current_lineages=False)
        # converted_lineage.orthospecies_index = self.current_orthospecies_index
        # converted_lineage.is_orthospecies = True
        # converted_lineage.speciation_completion_time = self.current_time
        # self.current_orthospecies_lineages.append(converted_lineage)
        # converted_lineage.node = original_lineage.node

        # original_lineage = self.rng.choice(self.current_incipient_species_lineages)
        # self.current_incipient_species_lineages.remove(original_lineage)
        # self.current_orthospecies_index += 1
        # converted_lineage = self._new_lineage(
        #         parent_lineage=original_lineage,
        #         orthospecies_index=self.current_orthospecies_index,
        #         is_orthospecies=True,
        #         add_to_current_lineages=False)
        # converted_lineage.orthospecies_index = self.current_orthospecies_index
        # converted_lineage.is_orthospecies = True
        # converted_lineage.speciation_completion_time = self.current_time
        # self.current_orthospecies_lineages.append(converted_lineage)
        # c1 = self._new_node(lineage=converted_lineage)
        # original_lineage.node.add_child(c1)

        # original_lineage = self.rng.choice(self.current_incipient_species_lineages)
        # self.current_incipient_species_lineages.remove(original_lineage)
        # self.current_orthospecies_index += 1
        # converted_lineage = self._new_lineage(
        #         parent_lineage=original_lineage,
        #         orthospecies_index=self.current_orthospecies_index,
        #         is_orthospecies=True,
        #         add_to_current_lineages=False)
        # converted_lineage.orthospecies_index = self.current_orthospecies_index
        # converted_lineage.is_orthospecies = True
        # converted_lineage.speciation_completion_time = self.current_time
        # self.current_orthospecies_lineages.append(converted_lineage)
        # p2 = original_lineage.node
        # c1 = self._new_node(lineage=converted_lineage)
        # c2 = self._new_node(lineage=original_lineage)
        # p2.add_child(c1)
        # p2.add_child(c2)

        # original_lineage = self.rng.choice(self.current_incipient_species_lineages)
        # self.current_incipient_species_lineages.remove(original_lineage)
        # self.current_orthospecies_index += 1
        # converted_lineage = original_lineage
        # converted_lineage.orthospecies_index = self.current_orthospecies_index
        # converted_lineage.is_orthospecies = True
        # converted_lineage.speciation_completion_time = self.current_time
        # self.current_orthospecies_lineages.append(converted_lineage)

        # lineage = self.rng.choice(self.current_incipient_species_lineages)
        # self.current_incipient_species_lineages.remove(lineage)
        # self.current_orthospecies_index += 1
        # converted_lineage = self._new_lineage(
        #         parent_lineage=lineage.parent_lineage,
        #         orthospecies_index=self.current_orthospecies_index,
        #         is_orthospecies =True)
        # converted_lineage.speciation_completion_time = self.current_time
        # self.current_orthospecies_lineages.append(converted_lineage)
        # converted_lineage.lineage_tree_node_history = list(lineage.lineage_tree_node_history)
        # old_node = lineage.node
        # c1 = self._new_node(lineage=converted_lineage)
        # c2 = self._new_node(lineage=lineage)
        # old_node.add_child(c1)
        # old_node.add_child(c2)

    def _process_orthospecies_extinction(self, tree):
        sp = self.rng.choice(self.current_orthospecies_lineages)
        sp.extinction_time = self.current_time
        self.current_orthospecies_lineages.remove(sp)
        self._make_lineage_extinct_on_phylogeny(tree, sp.node)

    def _process_incipient_species_extinction(self, tree):
        sp = self.rng.choice(self.current_incipient_species_lineages)
        sp.extinction_time = self.current_time
        self.current_incipient_species_lineages.remove(sp)
        self._make_lineage_extinct_on_phylogeny(tree, sp.node)

    def _make_lineage_extinct_on_phylogeny(self, tree, sp):
        if len(self.current_orthospecies_lineages) == 0 and len(self.current_incipient_species_lineages) == 0:
            raise TreeSimTotalExtinctionException()
        tree.prune_subtree(sp)

    def _new_lineage(self,
            parent_lineage,
            orthospecies_index,
            is_orthospecies,
            add_to_current_lineages=True):
        self.current_lineage_index += 1
        lineage_index = self.current_lineage_index
        speciation_initiation_time = self.current_time
        new_lineage = ProtractedSpeciationProcess.ProtractedSpeciationProcessLineage(
                index=lineage_index,
                parent_lineage=parent_lineage,
                speciation_initiation_time=speciation_initiation_time,
                is_orthospecies=is_orthospecies,
                orthospecies_index=orthospecies_index)
        if add_to_current_lineages:
            if is_orthospecies:
                self.current_orthospecies_lineages.append(new_lineage)
            else:
                self.current_incipient_species_lineages.append(new_lineage)
        return new_lineage

    def _new_node(self,
            lineage,
            ):
        node = self.node_factory()
        node.edge.length = 0.0
        node.protracted_speciation_model_lineage = lineage
        node.is_orthospeciation_event = False
        self.current_node_index += 1
        node.label = "{}.n{}".format(lineage.label, self.current_node_index)
        node.index = self.current_node_index
        node.annotations.add_new(name="lineage_index", value=lineage.index)
        node.annotations.add_new(name="lineage_label", value=lineage.label)
        node.annotations.add_new(name="speciation_initiation_time", value=lineage.speciation_initiation_time)
        lineage.node = node
        return node

    def _assemble_orthospecies_tree(self, taxon_namespace=None):
        lineage_set = set(self.current_incipient_species_lineages + self.current_orthospecies_lineages)
        sorted_lineages = sorted(lineage_set,
                key = lambda x: -x.speciation_initiation_time)
        self.lineage_to_orthospecies_tree_node_map = {}
        while sorted_lineages:
            lineage = sorted_lineages.pop(0)
            lineage_set.remove(lineage)
            parent_lineage = lineage.parent_lineage
            if parent_lineage is None:
                break
            if lineage.is_orthospecies:
                orthospecies_tree_node = self._require_orthospecies_tree_node(lineage=lineage)
                # try:
                #     orthospecies_tree_node = self.lineage_to_orthospecies_tree_node_map[lineage]
                # except KeyError:
                #     orthospecies_tree_node = dendropy.Node()
                #     orthospecies_tree_node.label = "L{}".format(lineage.index)
                #     orthospecies_tree_node.protracted_speciation_model_lineage = lineage
                #     self.lineage_to_orthospecies_tree_node_map[lineage] = orthospecies_tree_node
                if lineage.is_orthospecies:
                    parent_lineage.is_orthospecies = True
                # try:
                #     orthospecies_tree_parent_node = self.lineage_to_orthospecies_tree_node_map[parent_lineage]
                # except KeyError:
                #     orthospecies_tree_parent_node = dendropy.Node()
                #     orthospecies_tree_parent_node.label = "L{}".format(parent_lineage.index)
                #     orthospecies_tree_parent_node.protracted_speciation_model_lineage = parent_lineage
                #     self.lineage_to_orthospecies_tree_node_map[parent_lineage] = orthospecies_tree_parent_node
                orthospecies_tree_parent_node = self._require_orthospecies_tree_node(lineage=parent_lineage)
                orthospecies_tree_parent_node.add_child(orthospecies_tree_node)
                if parent_lineage not in lineage_set:
                    lineage_set.add(parent_lineage)
                    sorted_lineages = sorted(lineage_set,
                            key = lambda x: -x.speciation_initiation_time)

        # identify seed node
        seed_node = None
        for nd in self.lineage_to_orthospecies_tree_node_map.values():
            if nd.parent_node is None:
                seed_node = nd
                break
        if seed_node is None:
            raise ProcessFailedException()

        # create pruned tree
        orthospecies_tree = dendropy.Tree(taxon_namespace=taxon_namespace, seed_node=seed_node)
        orthospecies_tree.is_rooted = True

        # orthospecies_tree.suppress_unifurcations()
        # orthospecies_tree.set_edge_lengths_from_node_ages()

        ## assign ages to orthospecies tree
        for nd in orthospecies_tree.postorder_node_iter():
            if nd.is_leaf():
                nd.age = 0
            else:
                nd.age = self.current_time - min(ch.protracted_speciation_model_lineage.speciation_initiation_time for ch in nd.child_node_iter())
        orthospecies_tree.set_edge_lengths_from_node_ages()
        orthospecies_tree.suppress_unifurcations()
        return orthospecies_tree

    def _postprocess_psm_and_orthospecies_trees(self,
                        lineage_tree,
                        orthospecies_tree,
                        is_correlate_lineage_and_species_trees=False,
                        ):
        # lineage_tree.suppress_unifurcations()
        for nd in lineage_tree:
            nd.annotations.add_new(name="speciation_initiation_age", value=self.current_time - nd.annotations["speciation_initiation_time"].value)
            nd.annotations.add_new(name="lineage_orthospecies_index", value=nd.protracted_speciation_model_lineage.orthospecies_index)
            nd.label = "S{}.L{}.n{}".format(nd.protracted_speciation_model_lineage.orthospecies_index, nd.protracted_speciation_model_lineage.index, nd.index)
        if is_correlate_lineage_and_species_trees:
            raise NotImplementedError()
            # return self.correlate_lineage_and_species_trees(
            #         lineage_tree=lineage_tree,
            #         orthospecies_tree=orthospecies_tree)
        else:
            return lineage_tree, orthospecies_tree

    def _require_orthospecies_tree_node(self, lineage):
        try:
            node = self.lineage_to_orthospecies_tree_node_map[lineage]
        except KeyError:
            node = self.node_factory()
            node.label = "S{}".format(lineage.orthospecies_index)

            ## this might be misleading, as, technically, multiple lineages
            ## might be associated with this node; remove?
            node.protracted_speciation_model_lineage = lineage
            node.annotations.add_new(name="lineage_index", value=lineage.index)
            node.annotations.add_new(name="lineage_label", value=lineage.label)
            node.annotations.add_new(name="lineage_orthospecies_index", value=lineage.orthospecies_index)

            self.lineage_to_orthospecies_tree_node_map[lineage] = node
        return node

    def _debug_dump_lineages(self, lineages):
        sorted_lineages = sorted(lineages,
                key = lambda x: -x.speciation_initiation_time)
        for k in sorted_lineages:
            if k.parent_lineage is None:
                pi = "NA"
                pt = "NA"
            else:
                pi = k.parent_lineage.index
                pt = k.parent_lineage.is_orthospecies
            sys.stderr.write("{:10.5f} : {:4} ({}) => {} ({})\n".format(
                    k.speciation_initiation_time,
                    k.index,
                    k.is_orthospecies,
                    pi,
                    pt))
        sys.stderr.write("\n")