#!/usr/bin/env python
from peyotl.ott import create_pruned_and_taxonomy_for_tip_ott_ids
from peyotl.phylo.compat import compare_bits_as_splits, SplitComparison
from peyotl.utility import any_early_exit, get_logger
from peyotl.utility.str_util import reverse_dict
_LOG = get_logger(__name__)
def evaluate_tree_rooting(nexson, ott, tree_proxy):
    '''
    Returns None if the taxanomy contributes no information to the rooting decision
        (e.g. all of the tips are within one genus in the taxonomy)

    TODO: need to coordinate with Jim Allman and see if we can
    do this in a non- O(nm) manner (where n and m are the # of non-trivial edges in the phylo and taxo tree)
    putting barrier notes on the phylo tree would work...

    '''
    pruned_phylo, taxo_tree = create_pruned_and_taxonomy_for_tip_ott_ids(tree_proxy, ott)
    if taxo_tree is None: # this can happen if no otus are mapped
        return None
    has_taxo_groupings = any_early_exit(taxo_tree.root.child_iter(), lambda nd: not nd.is_leaf)
    if not has_taxo_groupings:
        return None
    has_phylo_groupings = any_early_exit(pruned_phylo.root.child_iter(), lambda nd: not nd.is_leaf)
    if not has_phylo_groupings:
        return None
    id2bit = pruned_phylo.add_bits4subtree_ids(None)
    taxo_tree.add_bits4subtree_ids(id2bit)
    assert taxo_tree.root.bits4subtree_ids == pruned_phylo.root.bits4subtree_ids
    taxo_nontriv_splits = taxo_tree.bits2internal_node
    taxon_mask = taxo_tree.root.bits4subtree_ids
    #_LOG.debug('taxo_nontriv_splits = {}'.format(taxo_nontriv_splits))
    # might want to copy this dict rather than modify in place..
    del taxo_nontriv_splits[taxon_mask] # root bitmask is trivial
    _LOG.debug('taxon_mask = {} (which is {} bits)'.format(bin(taxon_mask)[2:], len(bin(taxon_mask)) - 2))
    num_ids = len(id2bit)
    _LOG.debug('id2bit has length = {}'.format(len(id2bit)))
    #for checking tips of the phylogeny, it is nice to know which leaf OTUs attach
    #   at the base of the taxonomy (no other grouping)
    basal_taxo = set()
    basal_bits = 0
    for c in taxo_tree.root.child_iter():
        if c.is_leaf:
            basal_taxo.add(c._id)
            basal_bits |= id2bit[c._id]
    _LOG.debug('basal_bits = {}'.format(bin(basal_bits)[2:].zfill(num_ids)))
    _LOG.debug('# nontrivial taxo splits = {}'.format(len(taxo_nontriv_splits)))
    _EMPTY_SET = frozenset([])
    non_root_pp_preorder = [nd for nd in pruned_phylo.preorder_node_iter()][1:]
    curr_root_incompat_set = set()
    any_root_incompat_set = set()
    _taxo_node_id_set_cache = {_EMPTY_SET: _EMPTY_SET}
    for node in non_root_pp_preorder:
        edge = node.edge
        if node.is_leaf:
            edge._displays = None
            edge._inverted_displays = None
            b = id2bit[node._id]
            if node._id in basal_taxo:
                edge._not_inverted_incompat = _EMPTY_SET
                edge._inverted_incompat = _EMPTY_SET
                inv_mask = taxon_mask - b
                idisp = taxo_nontriv_splits.get(inv_mask)
                if idisp is not None:
                    edge._inverted_displays = idisp
            else:
                edge._not_inverted_incompat = _EMPTY_SET
                #TODO would be more efficient to jump to tip and walk back...
                b = id2bit[node._id]
                ii = set()
                for tb, tid in taxo_nontriv_splits.items():
                    if tb & b:
                        ii.add(tid)
                edge._inverted_incompat = _get_cached_set(ii, _taxo_node_id_set_cache)
                disp = taxo_nontriv_splits.get(b)
                if disp is not None:
                    edge._displays = disp
        else:
            #TODO this could be more efficient...
            b = node.bits4subtree_ids