#!/usr/bin/env python

"""Enum submodule for finding or iterating over topologies from
a defined tree space.

IN DEVELOPMENT

TODO
----
- Yule tree probs
- get n-maximally probable tree

"""

import numpy as np
from toytree.core import ToyTree


def get_unlabeled_trees(ntips: int) -> int:
    r"""Return unlabaled tree and its probability.

    Note
    ----
    In contrast to labeled trees, all unlabeled trees for ntips are 
    *not* equiprobable, since there are some topology shapes that can
    be formed more than others by swapping around the labels on a tree.
    ...
    """


def get_the_n_maximally_probable_tree(ntips: int) -> ToyTree:
    """Return the topology most probable under Yule model for ntips.

    Reference
    ---------
    - ...
    """


def get_yule_tree_probability(tree: ToyTree) -> float:
    """Return the probability of ... under a Yule model.

    The probability that a specific labeled history is generated under
    a Yule model is ... equation 3+> of Brown...

    Utility
    -------
    Under the Yule model two (and only two) lineages are randomly
    joined at a time until all lineages have been merged into a tree.
    For nlineages there is thus a space of possible Yule trees, within
    which not all tree topologies are equally likely. This function
    returns the probabilitiy of the input tree topology under the Yule
    model, given ntips.

    Reference
    ---------
    - Harding EF (1971) The probabilities of rooted tree-shapes generated by random bifurcation. Adv Appl Prob 3: 44–77.
    - Brown JKM (1994) Probabilities of evolutionary trees. Syst Biol 43: 78–91.
    - Aldous DJ (2001) Stochastic models and descriptive statistics for phylogenetic trees, from Yule to today. Stat Sci 16: 23–34.
    - Steel M, McKenzie A (2001) Properties of phylogenetic trees generated by Yule-type speciation models. Math Biosci 170: 91–112.
    """


def get_probability_gene_tree_matches_species_tree(internal_branch: float) -> float:
    """...

    TODO: use in the textbook.

    References
    ----------
    - Pamilo and Nei (1988)
    """

    # a quartet (rooted three taxon) tree has no anomalies b/c
    λ = internal_branch

    # probability the gene tree matches the species tree
    prob_true = 1 - (2 / 3) * np.exp(-λ)

    # probability the gene tree matches other 2 possible trees
    prob_false = (1 / 3) * np.exp(-λ)

    # no matter how small the internal branch, prob_true always > prob_false
    return prob_true > prob_false


if __name__ == "__main__":
    pass
