#! /usr/bin/env python3
from .taxonomy_util import Taxonomy
from .erlang import erlang
import math

class TaxTreeHelper:
    def __init__(self, cfg, tax_map, tax_tree=None):
        self.origin_taxonomy = tax_map
        self.cfg = cfg
        self.outgroup = None
        self.mf_rooted_tree = None
        self.bf_rooted_tree = None
        self.tax_tree = tax_tree
        self.bid_taxonomy_map = None
        self.ranks_set = set()
        if tax_tree:
            self.init_taxnode_map()
        else:
            self.name2taxnode = {}
    
    def set_mf_rooted_tree(self, rt):
        self.mf_rooted_tree = rt
        self.save_outgroup()
        self.bf_rooted_tree = None
        self.tax_tree = None
        self.bid_taxonomy_map = None
    
    def set_bf_unrooted_tree(self, ut):
        self.restore_rooting(ut)
        
    def get_outgroup(self):
        return self.outgroup

    def set_outgroup(self, outgr):
        self.outgroup = outgr

    def get_tax_tree(self):
        if not self.tax_tree:
            self.label_bf_tree_with_ranks()
        return self.tax_tree
    
    def get_bid_taxonomy_map(self, rebuild=False):
        self.get_tax_tree()
        if not self.bid_taxonomy_map or rebuild:
            self.build_bid_taxonomy_map()
        return self.bid_taxonomy_map

    def init_taxnode_map(self):
        self.name2taxnode = {}
        for leaf in self.tax_tree.iter_leaves():
            self.name2taxnode[leaf.name] = leaf
        
    def save_outgroup(self):
        rt = self.mf_rooted_tree
        
        # remove unifurcation at the root
        if len(rt.children) == 1:
            rt = rt.children[0]

        if len(rt.children) > 1:
            outgr = rt.children[0]    
            outgr_size = len(outgr.get_leaves())
            for child in rt.children:
                if child != outgr:
                    child_size = len(child.get_leaves())
                    if child_size < outgr_size:
                        outgr = child
                        outgr_size = child_size
        else:
            raise AssertionError("Invalid tree: unifurcation at the root node!")
        
        self.outgroup = outgr
    
    def restore_rooting(self, utree):
        outgr_leaves = self.outgroup.get_leaf_names()
        # check if outgroup consists of a single node - ETE considers it to be root, not leaf
        if not outgr_leaves:
            outgr_root = utree&outgr.name
        elif len(outgr_leaves) == 1:
            outgr_root = utree&outgr_leaves[0]
        else:
            # Even unrooted tree is "implicitely" rooted in ETE representation.
            # If this pseudo-rooting happens to be within the outgroup, it cause problems
            # in the get_common_ancestor() step (since common_ancestor = "root")
            # Workaround: explicitely root the tree outside from outgroup subtree
            for node in utree.iter_leaves():
                if not node.name in outgr_leaves:
                    tmp_root = node.up
                    if not utree == tmp_root:
                        utree.set_outgroup(tmp_root)
                        break
            
            outgr_root = utree.get_common_ancestor(outgr_leaves)

        # we could be so lucky that the RAxML tree is already correctly rooted :)
        if outgr_root != utree:
            utree.set_outgroup(outgr_root)

        self.bf_rooted_tree = utree

    def label_bf_tree_with_ranks(self):
        """labeling inner tree nodes with taxonomic ranks"""
        if not self.bf_rooted_tree:
            raise AssertionError("self.bf_rooted_tree is not set: TaxTreeHelper.set_bf_unrooted_tree() must be called before!")
            
        for node in self.bf_rooted_tree.traverse("postorder"):
            if node.is_leaf():
                seq_ranks = self.origin_taxonomy[node.name]
                rank_level = Taxonomy.lowest_assigned_rank_level(seq_ranks)
                node.add_feature("rank_level", rank_level)
                node.add_feature("ranks", seq_ranks)
                node.name += "__" + seq_ranks[rank_level]
            else:
                if len(node.children) != 2:
                    raise AssertionError("FATAL ERROR: tree is not bifurcating!")
                lchild = node.children[0]
                rchild = node.children[1]
                rank_level = min(lchild.rank_level, rchild.rank_level)
                while rank_level >= 0 and lchild.ranks[rank_level] != rchild.ranks[rank_level]:
                    rank_level -= 1
                node.add_feature("rank_level", rank_level)
                node_ranks = [Taxonomy.EMPTY_RANK] * max(len(lchild.ranks),len(rchild.ranks)) 
                if rank_level >= 0:
                    node_ranks[0:rank_level+1] = lchild.ranks[0:rank_level+1]
                    node.name = lchild.ranks[rank_level]
                else:
                    node.name = "Undefined"
                    if hasattr(node, "B"):
                        self.cfg.log.debug("INFO: empty taxonomic annotation for branch %s (child nodes have no common ranks)", node.B)
                
                node.add_feature("ranks", node_ranks)

        self.tax_tree = self.bf_rooted_tree
        self.init_taxnode_map()

    def build_bid_taxonomy_map(self):
        self.bid_taxonomy_map = {}
        self.ranks_set = set([])
        for node in self.tax_tree.traverse("postorder"):
            if not node.is_root() and hasattr(node, "B"):
                parent = node.up
                branch_rdiff = Taxonomy.lowest_assigned_rank_level(node.ranks) - Taxonomy.lowest_assigned_rank_level(parent.ranks)
                branch_rank_id = Taxonomy.get_rank_uid(node.ranks)
                branch_len = node.dist
                self.bid_taxonomy_map[node.B] = (branch_rank_id, branch_rdiff, branch_len)
                self.ranks_set.add(branch_rank_id)
#                if self.cfg.debug:
#                  print node.ranks, parent.ranks, branch_diff

    def get_seq_ranks_from_tree(self, seq_name):
        if seq_name not in self.name2taxnode:
            errmsg = "FATAL ERROR: Sequence %s is not found in the taxonomic tree!" % seq_name
            self.cfg.exit_fatal_error(errmsg)

        seq_node = self.name2taxnode[seq_name]
        ranks = Taxonomy.split_rank_uid(seq_node.up.name)
        return ranks

    def strip_missing_ranks(self, ranks):
        rank_level = len(ranks)
        while not Taxonomy.get_rank_uid(ranks[0:rank_level]) in self.ranks_set and rank_level > 0:
            rank_level -= 1
        
        return ranks[0:rank_level]   
    
class TaxClassifyHelper:
    def __init__(self, cfg, bid_taxonomy_map, sp_rate = 0., node_height = []):
        self.cfg = cfg
        self.bid_taxonomy_map = bid_taxonomy_map
        self.sp_rate = sp_rate
        self.node_height = node_height
        self.erlang = erlang()
        # hardcoded for now
        self.parent_lhw_coeff = 0.49

    def classify_seq(self, edges, minlw = None):
        if not minlw:
            minlw = self.cfg.min_lhw

        edges = self.erlang_filter(edges)
        if len(edges) > 0:
            if self.cfg.taxassign_method == "1":
                ranks, lws = self.assign_taxonomy_maxsum(edges, minlw)
            else:
                ranks, lws = self.assign_taxonomy_maxlh(edges)
            return ranks, lws
        else:
            return [], []      
            
    def erlang_filter(self, edges):
        if self.cfg.brlen_pv == 0.:
            return edges
            
        newedges = []
        for edge in edges:
            edge_nr = str(edge[0])
            pendant_length = edge[4]
            pv = self.erlang.one_tail_test(rate = self.sp_rate, k = int(self.node_height[edge_nr]), x = pendant_length)
            if pv >= self.cfg.brlen_pv:
                newedges.append(edge)
#            else:
#                self.cfg.log.debug("Edge ignored: [%s, %f], p = %.12f", edge_nr, pendant_length, pv)
        
        if len(newedges) == 0:
            return newedges
        
        # adjust likelihood weights -> is there a better way ???        
        sum_lh = 0
        max_lh = float(newedges[0][1])
        for edge in newedges:
            lh = float(edge[1])
            sum_lh += math.exp(lh - max_lh)
                        
        for edge in newedges:
            lh = float(edge[1])
            edge[2] = math.exp(lh - max_lh) / sum_lh

        return newedges

    # "all or none" filter: return empty set iff *all* brlens are below the threshold
    def erlang_filter2(self, edges):
        if self.cfg.brlen_pv == 0.:
            return edges
            
        for edge in edges:
            edge_nr = str(edge[0])
            pendant_length = edge[4]
            pv = self.erlang.one_tail_test(rate = self.sp_rate, k = int(self.node_height[edge_nr]), x = pendant_length)
            if pv >= self.cfg.brlen_pv:
                return edges
                
        return []
     
    def get_branch_ranks(self, br_id):
        br_rec = self.bid_taxonomy_map[br_id]
        br_rank_id = br_rec[0]
        ranks = Taxonomy.split_rank_uid(br_rank_id)            
        return ranks
    
    def assign_taxonomy_maxlh(self, edges):
        #Calculate the sum of likelihood weight for each rank
        taxonmy_sumlw_map = {}
        for edge in edges:
            edge_nr = str(edge[0])
            lw = edge[2]
            taxranks = self.get_branch_ranks(edge_nr)            
            for rank in taxranks:
                if rank == "-":
                    taxonmy_sumlw_map[rank] = -1
                elif rank in taxonmy_sumlw_map:
                    oldlw = taxonmy_sumlw_map[rank]
                    taxonmy_sumlw_map[rank] = oldlw + lw
                else:
                    taxonmy_sumlw_map[rank] = lw
        
        #Assignment using the max likelihood placement
        ml_edge = edges[0]
        edge_nr = str(ml_edge[0])
        maxlw = ml_edge[2]
        ml_ranks = self.get_branch_ranks(edge_nr)
        ml_ranks_copy = []
        for rk in ml_ranks:
            ml_ranks_copy.append(rk)
        lws = []
        cnt = 0
        for rank in ml_ranks:
            lw = taxonmy_sumlw_map[rank]
            if lw > 1.0:
                lw = 1.0
            lws.append(lw)
            if rank == "-" and cnt > 0 :                
                for edge in edges[1:]:
                    edge_nr = str(edge[0])
                    taxonomy = self.get_branch_ranks(edge_nr)
                    newrank = taxonomy[cnt]
                    newlw = taxonmy_sumlw_map[newrank]
                    higherrank_old = ml_ranks[cnt -1]
                    higherrank_new = taxonomy[cnt -1]
                    if higherrank_old == higherrank_new and newrank!="-":
                        ml_ranks_copy[cnt] = newrank
                        lws[cnt] = newlw
            cnt = cnt + 1
            
        return ml_ranks_copy, lws

    def assign_taxonomy_maxsum(self, edges, minlw):
        """this function sums up all LH-weights for each rank and takes the rank with the max. sum """
        # in EPA result, each placement(=branch) has a "weight"
        # since we are interested in taxonomic placement, we do not care about branch vs. branch comparisons,
        # but only consider rank vs. rank (e. g. G1 S1 vs. G1 S2 vs. G1)
        # Thus we accumulate weights for each rank, there are to measures:
        # "own" weight  = sum of weight of all placements EXACTLY to this rank (e.g. for G1: G1 only)
        # "total" rank  = own rank + own rank of all children (for G1: G1 or G1 S1 or G1 S2)
        rw_own = {}
        rw_total = {}
        
        ranks = [Taxonomy.EMPTY_RANK]
        
        for edge in edges:
            br_id = str(edge[0])
            lweight = edge[2]
            lowest_rank = None
            lowest_rank_lvl = None

            if lweight == 0.:
                continue

            # accumulate weight for the current sequence                
            br_rank_id, rdiff, brlen = self.bid_taxonomy_map[br_id]
            ranks = Taxonomy.split_rank_uid(br_rank_id)
            for i in range(len(ranks)):
                rank = ranks[i]
                rank_id = Taxonomy.get_rank_uid(ranks, i)
                if rank != Taxonomy.EMPTY_RANK:
                    rw_total[rank_id] = rw_total.get(rank_id, 0) + lweight
                    lowest_rank_lvl = i
                    lowest_rank = rank_id
                else:
                    break

            if lowest_rank:
                if rdiff > 0:
                  # if ranks of 'upper' and 'lower' adjacent nodes of a branch are non-equal, split LHW among them
                  parent_rank = Taxonomy.get_rank_uid(ranks, lowest_rank_lvl - rdiff)
                  rw_own[lowest_rank] = rw_own.get(lowest_rank, 0) + lweight * (1 - self.parent_lhw_coeff)
                  rw_own[parent_rank] = rw_own.get(parent_rank, 0) + lweight * self.parent_lhw_coeff
                  # correct total lhw for all levels between "parent" and "lowest"
                  # NOTE: all those intermediate ranks are in fact indistinguishable, e.g. a family which contains a single genus
                  for r in range(rdiff):
                    interim_rank = Taxonomy.get_rank_uid(ranks, lowest_rank_lvl - r)
                    rw_total[interim_rank] = rw_total.get(interim_rank, 0) - lweight * self.parent_lhw_coeff
                else:
                  rw_own[lowest_rank] = rw_own.get(lowest_rank, 0) + lweight
#            else:
#                self.cfg.log.debug("WARNING: no annotation for branch %s", br_id)
            
        # if all branches have empty ranks only, just return this placement
        if len(rw_total) == 0:
            return ranks, [1.] * len(ranks)
        
        # we assign the sequence to a rank, which has the max "own" weight AND 
        # whose "total" weight is greater than a confidence threshold
        max_rw = 0.
        ass_rank_id = None
        for r in rw_own.keys():
            if rw_own[r] > max_rw and rw_total[r] >= minlw:
                ass_rank_id = r
                max_rw = rw_own[r] 
        if not ass_rank_id:
            ass_rank_id = max(iter(rw_total.keys()), key=(lambda key: rw_total[key]))

        a_ranks = Taxonomy.split_rank_uid(ass_rank_id)
        
        # "total" weight is considered as confidence value for now
        a_conf = [0.] * len(a_ranks)
        for i in range(len(a_conf)):
            rank = a_ranks[i]
            if rank != Taxonomy.EMPTY_RANK:
                rank_id = Taxonomy.get_rank_uid(a_ranks, i)
                a_conf[i] = rw_total[rank_id]

        return a_ranks, a_conf
    
