#! /usr/bin/env python
# -*- coding: utf-8 -*-

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

"""
Classes and Methods for working with tree reconciliation, fitting, embedding,
contained/containing etc.
"""

import dendropy
from dendropy.model import coalescent

class ContainingTree(dendropy.Tree):
    """
    A "containing tree" is a (usually rooted) tree data structure within which
    other trees are "contained". For example, species trees and their contained
    gene trees; host trees and their contained parasite trees; biogeographical
    "area" trees and their contained species or taxon trees.
    """

    def __init__(self,
            containing_tree,
            contained_taxon_namespace,
            contained_to_containing_taxon_map,
            contained_trees=None,
            fit_containing_edge_lengths=True,
            collapse_empty_edges=True,
            ultrametricity_precision=False,
            ignore_root_deep_coalescences=True,
            **kwargs):
        """
        __init__ converts ``self`` to ContainingTree class, embedding the trees
        given in the list, ``contained_trees.``


        Mandatory Arguments:

            ``containing_tree``
                A |Tree| or |Tree|-like object that describes the topological
                constraints or conditions of the containing tree (e.g., species,
                host, or biogeographical area trees).

            ``contained_taxon_namespace``
                A |TaxonNamespace| object that will be used to manage the taxa of
                the contained trees.

            ``contained_to_containing_taxon_map``
                A |TaxonNamespaceMapping| object mapping |Taxon| objects in the
                contained |TaxonNamespace| to corresponding |Taxon| objects in the
                containing tree.

        Optional Arguments:

            ``contained_trees``
                An iterable container of |Tree| or |Tree|-like objects that
                will be contained into ``containing_tree``; e.g. gene or
                parasite trees.

            ``fit_containing_edge_lengths``
                If |True| [default], then the branch lengths of
                ``containing_tree`` will be adjusted to fit the contained tree
                as they are added. Otherwise, the containing tree edge lengths
                will not be changed.

            ``collapse_empty_edges``
                If |True| [default], after edge lengths are adjusted,
                zero-length branches will be collapsed.

            ``ultrametricity_precision``
                If |False| [default], then trees will not be checked for
                ultrametricity. Otherwise this is the threshold within which
                all node to tip distances for sister nodes must be equal.

            ``ignore_root_deep_coalescences``
                If |True| [default], then deep coalescences in the root will
                not be counted.

        Other Keyword Arguments: Will be passed to Tree().

    """
        if "taxon_namespace" not in kwargs:
            kwargs["taxon_namespace"] = containing_tree.taxon_namespace
        dendropy.Tree.__init__(self,
                containing_tree,
                taxon_namespace=containing_tree.taxon_namespace)
        self.original_tree = containing_tree
        for edge in self.postorder_edge_iter():
            edge.head_contained_edges = {}
            edge.tail_contained_edges = {}
            edge.containing_taxa = set()
            edge.contained_taxa = set()
        self._contained_taxon_namespace = contained_taxon_namespace
        self._contained_to_containing_taxon_map = None
        self._contained_trees = None
        self._set_contained_to_containing_taxon_map(contained_to_containing_taxon_map)
        self.fit_containing_edge_lengths = fit_containing_edge_lengths
        self.collapse_empty_edges = collapse_empty_edges
        self.ultrametricity_precision = ultrametricity_precision
        self.ignore_root_deep_coalescences = ignore_root_deep_coalescences
        if contained_trees:
            self._set_contained_trees(contained_trees)
        if self.contained_trees:
            self.rebuild(rebuild_taxa=False)

    def _set_contained_taxon_namespace(self, taxon_namespace):
        self._contained_taxon_namespace = taxon_namespace

    def _get_contained_taxon_namespace(self):
        if self._contained_taxon_namespace is None:
            self._contained_taxon_namespace = dendropy.TaxonNamespace()
        return self._contained_taxon_namespace

    contained_taxon_namespace = property(_get_contained_taxon_namespace)

    def _set_contained_to_containing_taxon_map(self, contained_to_containing_taxon_map):
        """
        Sets mapping of |Taxon| objects of the genes/parasite/etc. to that of
        the population/species/host/etc.
        Creates mapping (e.g., species to genes) and decorates edges of self
        with sets of both containing |Taxon| objects and the contained
        |Taxon| objects that map to them.
        """
        if isinstance(contained_to_containing_taxon_map, dendropy.TaxonNamespaceMapping):
            if self._contained_taxon_namespace is not contained_to_containing_taxon_map.domain_taxon_namespace:
                raise ValueError("Domain TaxonNamespace of TaxonNamespaceMapping ('domain_taxon_namespace') not the same as 'contained_taxon_namespace' TaxonNamespace")
            self._contained_to_containing_taxon_map = contained_to_containing_taxon_map
        else:
            self._contained_to_containing_taxon_map = dendropy.TaxonNamespaceMapping(
                    mapping_dict=contained_to_containing_taxon_map,
                    domain_taxon_namespace=self.contained_taxon_namespace,
                    range_taxon_namespace=self.taxon_namespace)
        self.build_edge_taxa_sets()

    def _get_contained_to_containing_taxon_map(self):
        return self._contained_to_containing_taxon_map

    contained_to_containing_taxon_map = property(_get_contained_to_containing_taxon_map)

    def _set_contained_trees(self, trees):
        if hasattr(trees, 'taxon_namespace'):
            if self._contained_taxon_namespace is None:
                self._contained_taxon_namespace = trees.taxon_namespace
            elif self._contained_taxon_namespace is not trees.taxon_namespace:
                raise ValueError("'contained_taxon_namespace' of ContainingTree is not the same TaxonNamespace object of 'contained_trees'")
        self._contained_trees = dendropy.TreeList(trees, taxon_namespace=self._contained_taxon_namespace)
        if self._contained_taxon_namespace is None:
            self._contained_taxon_namespace = self._contained_trees.taxon_namespace

    def _get_contained_trees(self):
        if self._contained_trees is None:
            self._contained_trees = dendropy.TreeList(taxon_namespace=self._contained_taxon_namespace)
        return self._contained_trees

    contained_trees = property(_get_contained_trees)

    def _get_containing_to_contained_taxa_map(self):
        return self._contained_to_containing_taxon_map.reverse

    containing_to_contained_taxa_map = property(_get_containing_to_contained_taxa_map)

    def clear(self):
        """
        Clears all contained trees and mapped edges.
        """
        self.contained_trees = dendropy.TreeList(taxon_namespace=self._contained_to_containing_taxon_map.domain_taxa)
        self.clear_contained_edges()

    def clear_contained_edges(self):
        """
        Clears all contained mapped edges.
        """
        for edge in self.postorder_edge_iter():
            edge.head_contained_edges = {}
            edge.tail_contained_edges = {}

    def fit_edge_lengths(self, contained_trees):
        """
        Recalculate node ages / edge lengths of containing tree to accomodate
        contained trees.
        """

        # set the ages
        for node in self.postorder_node_iter():
            if node.is_internal():
                disjunct_leaf_set_list_split_bitmasks = []
                for i in node.child_nodes():
                    disjunct_leaf_set_list_split_bitmasks.append(self.taxon_namespace.taxa_bitmask(taxa=i.edge.containing_taxa))
                min_age = float('inf')
                for et in contained_trees:
                    min_age = self._find_youngest_intergroup_age(et, disjunct_leaf_set_list_split_bitmasks, min_age)
                node.age = max( [min_age] + [cn.age for cn in node.child_nodes()] )
            else:
                node.age = 0

        # set the corresponding edge lengths
        self.set_edge_lengths_from_node_ages()

        # collapse 0-length branches
        if self.collapse_empty_edges:
           self.collapse_unweighted_edges()

    def rebuild(self, rebuild_taxa=True):
        """
        Recalculate edge taxa sets, node ages / edge lengths of containing
        tree, and embed edges of contained trees.
        """
        if rebuild_taxa:
            self.build_edge_taxa_sets()
        if self.fit_containing_edge_lengths:
            self.fit_edge_lengths(self.contained_trees)
        self.clear_contained_edges()
        for et in self.contained_trees:
            self.embed_tree(et)

    def embed_tree(self, contained_tree):
        """
        Map edges of contained tree into containing tree (i.e., self).
        """
        if self.seed_node.age is None:
            self.calc_node_ages(ultrametricity_precision=self.ultrametricity_precision)
        if contained_tree not in self.contained_trees:
            self.contained_trees.append(contained_tree)
        if self.fit_containing_edge_lengths:
            self.fit_edge_lengths(self.contained_trees)
        if contained_tree.seed_node.age is None:
            contained_tree.calc_node_ages(ultrametricity_precision=self.ultrametricity_precision)
        contained_leaves = contained_tree.leaf_nodes()
        taxon_to_contained = {}
        for nd in contained_leaves:
            containing_taxon = self.contained_to_containing_taxon_map[nd.taxon]
            x = taxon_to_contained.setdefault(containing_taxon, set())
            x.add(nd.edge)
        for containing_edge in self.postorder_edge_iter():
            if containing_edge.is_terminal():
                containing_edge.head_contained_edges[contained_tree] = taxon_to_contained[containing_edge.head_node.taxon]
            else:
                containing_edge.head_contained_edges[contained_tree] = set()
                for nd in containing_edge.head_node.child_nodes():
                    containing_edge.head_contained_edges[contained_tree].update(nd.edge.tail_contained_edges[contained_tree])

            if containing_edge.tail_node is None:
                if containing_edge.length is not None:
                    target_age =  containing_edge.head_node.age + containing_edge.length
                else:
                    # assume all coalesce?
                    containing_edge.tail_contained_edges[contained_tree] = set([contained_tree.seed_node.edge])
                    continue
            else:
                target_age = containing_edge.tail_node.age

            containing_edge.tail_contained_edges[contained_tree] = set()
            for contained_edge in containing_edge.head_contained_edges[contained_tree]:
                if contained_edge.tail_node is not None:
                    remaining = target_age - contained_edge.tail_node.age
                elif contained_edge.length is not None:
                    remaining = target_age - (contained_edge.head_node.age + contained_edge.length)
                else:
                    continue
                while remaining > 0:
                    if contained_edge.tail_node is not None:
                        contained_edge = contained_edge.tail_node.edge
                    else:
                        if contained_edge.length is not None and (remaining - contained_edge.length) <= 0:
                            contained_edge = None
                            remaining = 0
                            break
                        else:
                            remaining = 0
                            break
                    if contained_edge and remaining > 0:
                        remaining -= contained_edge.length
                if contained_edge is not None:
                    containing_edge.tail_contained_edges[contained_tree].add(contained_edge)

    def build_edge_taxa_sets(self):
        """
        Rebuilds sets of containing and corresponding contained taxa at each
        edge.
        """
        for edge in self.postorder_edge_iter():
            if edge.is_terminal():
                edge.containing_taxa = set([edge.head_node.taxon])
            else:
                edge.containing_taxa = set()
                for i in edge.head_node.child_nodes():
                    edge.containing_taxa.update(i.edge.containing_taxa)
            edge.contained_taxa = set()
            for t in edge.containing_taxa:
                edge.contained_taxa.update(self.containing_to_contained_taxa_map[t])

    def num_deep_coalescences(self):
        """
        Returns total number of deep coalescences of the contained trees.
        """
        return sum(self.deep_coalescences().values())

    def deep_coalescences(self):
        """
        Returns dictionary where the contained trees are keys, and the number of
        deep coalescences corresponding to the tree are values.
        """
        dc = {}
        for tree in self.contained_trees:
            for edge in self.postorder_edge_iter():
                if edge.tail_node is None and self.ignore_root_deep_coalescences:
                    continue
                try:
                    dc[tree] += len(edge.tail_contained_edges[tree]) - 1
                except KeyError:
                    dc[tree] = len(edge.tail_contained_edges[tree]) - 1
        return dc

    def embed_contained_kingman(self,
            edge_pop_size_attr='pop_size',
            default_pop_size=1,
            label=None,
            rng=None,
            use_expected_tmrca=False):
        """
        Simulates, *embeds*, and returns a "censored" (Kingman) neutral coalescence tree
        conditional on self.

            ``rng``
                Random number generator to use. If |None|, the default will
                be used.

            ``edge_pop_size_attr``
                Name of attribute of self's edges that specify the population
                size. If this attribute does not exist, then the population
                size is taken to be 1.

        Note that all edge-associated taxon sets must be up-to-date (otherwise,
        ``build_edge_taxa_sets()`` should be called).
        """
        et = self.simulate_contained_kingman(
                edge_pop_size_attr=edge_pop_size_attr,
                default_pop_size=default_pop_size,
                label=label,
                rng=rng,
                use_expected_tmrca=use_expected_tmrca)
        self.embed_tree(et)
        return et

    def simulate_contained_kingman(self,
            edge_pop_size_attr='pop_size',
            default_pop_size=1,
            label=None,
            rng=None,
            use_expected_tmrca=False):
        """
        Simulates and returns a "censored" (Kingman) neutral coalescence tree
        conditional on self.

            ``rng``
                Random number generator to use. If |None|, the default will
                be used.

            ``edge_pop_size_attr``
                Name of attribute of self's edges that specify the population
                size. If this attribute does not exist, then the population
                size is taken to be 1.

        Each coalescent tree terminal node will have a ``container_tree_node``
        attribute added that references the node on the container tree from
        which it was sampled.

        Note that all edge-associated taxon sets must be up-to-date (otherwise,
        ``build_edge_taxa_sets()`` should be called), and that the tree
        is *not* added to the set of contained trees. For the latter, call
        ``embed_contained_kingman``.
        """

        # Dictionary that maps nodes of containing tree to list of
        # corresponding nodes on gene tree, initially populated with leaf
        # nodes.
        contained_nodes = {}
        for nd in self.leaf_node_iter():
            contained_nodes[nd] = []
            for gt in nd.edge.contained_taxa:
                gn = dendropy.Node(taxon=gt)
                contained_nodes[nd].append(gn)
                gn.container_tree_node = nd

        # Generate the tree structure
        for edge in self.postorder_edge_iter():
            if edge.head_node.parent_node is None:
                # root: run unconstrained coalescence until just one gene node
                # remaining
                if hasattr(edge, edge_pop_size_attr):
                    pop_size = getattr(edge, edge_pop_size_attr)
                else:
                    pop_size = default_pop_size
                if len(contained_nodes[edge.head_node]) > 1:
                    final = coalescent.coalesce_nodes(nodes=contained_nodes[edge.head_node],
                            pop_size=pop_size,
                            period=None,
                            rng=rng,
                            use_expected_tmrca=use_expected_tmrca)
                else:
                    final = contained_nodes[edge.head_node]
            else:
                # run until next coalescence event, as determined by this edge
                # size.
                if hasattr(edge, edge_pop_size_attr):
                    pop_size = getattr(edge, edge_pop_size_attr)
                else:
                    pop_size = default_pop_size
                remaining = coalescent.coalesce_nodes(nodes=contained_nodes[edge.head_node],
                        pop_size=pop_size,
                        period=edge.length,
                        rng=rng,
                        use_expected_tmrca=use_expected_tmrca)
                try:
                    contained_nodes[edge.tail_node].extend(remaining)
                except KeyError:
                    contained_nodes[edge.tail_node] = remaining

        # Create and return the full tree
        contained_tree = dendropy.Tree(taxon_namespace=self.contained_taxon_namespace, label=label)
        contained_tree.seed_node = final[0]
        contained_tree.is_rooted = True
        return contained_tree

    def _find_youngest_intergroup_age(self, contained_tree, disjunct_leaf_set_list_split_bitmasks, starting_min_age=None):
        """
        Find the age of the youngest MRCA of disjunct leaf sets.
        """
        if starting_min_age is None:
            starting_min_age = float('inf')
        if contained_tree.seed_node.age is None:
            contained_tree.calc_node_ages(ultrametricity_precision=self.ultrametricity_precision)
        for nd in contained_tree.ageorder_node_iter(include_leaves=False):
            if nd.age > starting_min_age:
                break
            prev_intersections = False
            for bm in disjunct_leaf_set_list_split_bitmasks:
                if bm & nd.edge.split_bitmask:
                    if prev_intersections:
                        return nd.age
                    prev_intersections = True
        return starting_min_age

    def write_as_mesquite(self, out, **kwargs):
        """
        For debugging purposes, write out a Mesquite-format file.
        """
        from dendropy.dataio import nexuswriter
        nw = nexuswriter.NexusWriter(**kwargs)
        nw.is_write_block_titles = True
        out.write("#NEXUS\n\n")
        nw._write_taxa_block(out, self.taxon_namespace)
        out.write('\n')
        nw._write_taxa_block(out, self.contained_trees.taxon_namespace)
        if self.contained_trees.taxon_namespace.label:
            domain_title = self.contained_trees.taxon_namespace.label
        else:
            domain_title = self.contained_trees.taxon_namespace.oid
        contained_taxon_namespace = self.contained_trees.taxon_namespace
        contained_label = self.contained_trees.label
        out.write('\n')
        self._contained_to_containing_taxon_map.write_mesquite_association_block(out)
        out.write('\n')
        nw._write_trees_block(out, dendropy.TreeList([self], taxon_namespace=self.taxon_namespace))
        out.write('\n')
        nw._write_trees_block(out, dendropy.TreeList(self.contained_trees, taxon_namespace=contained_taxon_namespace, label=contained_label))
        out.write('\n')

def reconciliation_discordance(gene_tree, species_tree):
    """
    Given two trees (with splits encoded), this returns the number of gene
    duplications implied by the gene tree reconciled on the species tree, based
    on the algorithm described here:

        Goodman, M. J. Czelnusiniak, G. W. Moore, A. E. Romero-Herrera, and
        G. Matsuda. 1979. Fitting the gene lineage into its species lineage,
        a parsimony strategy illustrated by cladograms constructed from globin
        sequences. Syst. Zool. 19: 99-113.

        Maddison, W. P. 1997. Gene trees in species trees. Syst. Biol. 46:
        523-536.

    This function requires that the gene tree and species tree *have the same
    leaf set*. Note that for correct results,

        (a) trees must be rooted (i.e., is_rooted = True)
        (b) split masks must have been added as rooted (i.e., when
            encode_splits was called, is_rooted must have been set to True)

    """
    taxa_mask = species_tree.taxon_namespace.all_taxa_bitmask()
    species_node_gene_nodes = {}
    gene_node_species_nodes = {}
    for gnd in gene_tree.postorder_node_iter():
        gn_children = gnd.child_nodes()
        if len(gn_children) > 0:
            ssplit = 0
            for gn_child in gn_children:
                ssplit = ssplit | gene_node_species_nodes[gn_child].edge.leafset_bitmask
            sanc = species_tree.mrca(start_node=species_tree.seed_node, leafset_bitmask=ssplit)
            gene_node_species_nodes[gnd] = sanc
            if sanc not in species_node_gene_nodes:
                species_node_gene_nodes[sanc] = []
            species_node_gene_nodes[sanc].append(gnd)
        else:
            gene_node_species_nodes[gnd] = species_tree.find_node(lambda x : x.taxon == gnd.taxon)
    contained_gene_lineages = {}
    for snd in species_tree.postorder_node_iter():
        if snd in species_node_gene_nodes:
            for gnd in species_node_gene_nodes[snd]:
                for gnd_child in gnd.child_nodes():
                    sanc = gene_node_species_nodes[gnd_child]
                    p = sanc
                    while p is not None and p != snd:
                        if p.edge not in contained_gene_lineages:
                            contained_gene_lineages[p.edge] = 0
                        contained_gene_lineages[p.edge] += 1
                        p = p.parent_node

    dc = 0
    for v in contained_gene_lineages.values():
        dc += v - 1
    return dc

def monophyletic_partition_discordance(tree, taxon_namespace_partition):
    """
    Returns the number of deep coalescences on tree ``tree`` that would result
    if the taxa in ``tax_sets`` formed K mutually-exclusive monophyletic groups,
    where K = len(tax_sets)
    ``taxon_namespace_partition`` == TaxonNamespacePartition
    """

    tax_sets = taxon_namespace_partition.subsets()

    # from dendropy.model import parsimony
    # taxon_state_sets_map = {}
    # assert tree.taxon_namespace is taxon_namespace_partition.taxon_namespace
    # for taxon in tree.taxon_namespace:
    #     taxon_state_sets_map[taxon] = [0 for i in range(len(tax_sets))]
    # for idx, ts in enumerate(tax_sets):
    #     for taxon in ts:
    #         taxon_state_sets_map[taxon][idx] = 1
    # for taxon in tree.taxon_namespace:
    #     taxon_state_sets_map[taxon] = [set([i]) for i in taxon_state_sets_map[taxon]]
    # return parsimony.fitch_down_pass(
    #         postorder_nodes=tree.postorder_node_iter(),
    #         taxon_state_sets_map=taxon_state_sets_map
    #         )

    dc_tree = dendropy.Tree()
    dc_tree.taxon_namespace = dendropy.TaxonNamespace()
    for t in range(len(tax_sets)):
        dc_tree.taxon_namespace.add_taxon(dendropy.Taxon(label=str(t)))
    def _get_dc_taxon(nd):
        for idx, tax_set in enumerate(tax_sets):
            if nd.taxon in tax_set:
                return dc_tree.taxon_namespace[idx]
        assert "taxon not found in partition: '%s'" % nd.taxon.label
    src_dc_map = {}
    for snd in tree.postorder_node_iter():
        nnd = dendropy.Node()
        src_dc_map[snd] = nnd
        children = snd.child_nodes()
        if len(children) == 0:
            nnd.taxon = _get_dc_taxon(snd)
        else:
            taxa_set = []
            for cnd in children:
                dc_node = src_dc_map[cnd]
                if len(dc_node.child_nodes()) > 1:
                    nnd.add_child(dc_node)
                else:
                    ctax = dc_node.taxon
                    if ctax is not None and ctax not in taxa_set:
                        taxa_set.append(ctax)
                    del src_dc_map[cnd]
            if len(taxa_set) > 1:
                for t in taxa_set:
                    cnd = dendropy.Node()
                    cnd.taxon = t
                    nnd.add_child(cnd)
            else:
                if len(nnd.child_nodes()) == 0:
                    nnd.taxon = taxa_set[0]
                elif len(taxa_set) == 1:
                    cnd = dendropy.Node()
                    cnd.taxon = taxa_set[0]
                    nnd.add_child(cnd)
    dc_tree.seed_node = nnd
    return len(dc_tree.leaf_nodes()) - len(tax_sets)

