#!/usr/bin/env python3

import sys
import re
from .ete2 import Tree

class TaxCode:
    UNI_TAX_RANKS = {
         1: ("Domain", "d__"),
         2: ("Kingdom", "k__"),
         3: ("Subkingdom", "a__"),
         4: ("Phylum", "p__"),
         5: ("Subphylum", "a__"),
         6: ("Class", "c__"),
         7: ("Subclass", "x__"),
         8: ("Superorder", "e__"),
         9: ("Order", "o__"),
         10: ("Suborder", "h__"),
         11: ("Infraorder", "i__"),
         12: ("Superfamily", "j__"),
         13: ("Epifamily", "l__"),
         14: ("Family", "f__"), 
         15: ("Subfamily", "m__"),
         16: ("Infrafamily", "n__"),
         17: ("Tribe", "t__"),
         18: ("Subtribe", "u__"),
         19: ("Infratribe", "v__"), 
         20: ("Genus", "g__"),
         21: ("Species", "s__"),
         22: ("Subspecies", "b__"),
         23: ("Strain", "r__"),
         24: ("Isolate", "q__")
        }        
    UNI_TAX_LEVELS = len(UNI_TAX_RANKS)
    
    # ranks of the standard 7-levels taxonomy
    # kingdom, phylum, class, order, family, genus, species
    STD_RANKS = [2, 4, 6, 9, 14, 20, 21]
    
    BAC_TAX_CODE = {
         2: ((), ("bacteria", "archaea")),   # kingdom
         4: ((), ()),                        # phylum
         6: ((), ()),                        # class
         7: (("idae"), ()),                  # subclass
         9: (("ales"), ()),                  # order
         10: (("ineae"), ()),                 # suborder
         14: (("aceae"), ()),                # family
         15: (("oideae"), ()),               # subfamily
         20: ((), ()),                       # genus
         21: ((), ()),                       # species
         22: ((), ()),                       # subspecies
         23: ((), ()),                       # strain
         24: ((), ())                        # isolate
    }

    BOT_TAX_CODE = {
         1: ((), ("eukaryota")),                                # domain
         2: ((), ("plantae", "algae", "fungi")),                # kingdom
         3: ((), ("dikarya")),                                  # subkingdom
         4: (("phyta", "phycota", "mycota"), ()),               # phylum
         5: (("phytina", "phycotina", "mycotina"), ()),         # subphylum
         6: (("opsida", "phyceae", "mycetes"), ()),             # class
         7: (("idae", "phycidae", "mycetidae"), ()),            # subclass
         8: (("anae"), ()),                                     # superorder
         9: (("ales"), ()),                                     # order
         10: (("ineae"), ()),                                   # suborder
         11: (("aria"), ()),                                    # infraorder
         12: (("acea"), ()),                                    # superfamily
         14: (("aceae"), ()),                                   # family
         15: (("oideae"), ()),                                  # subfamily
         17: (("eae"), ()),                                     # tribe
         18: (("inae"), ()),                                    # subtribe
         20: ((), ()),                                          # genus
         21: ((), ()),                                          # species
         22: ((), ()),                                          # subspecies
         23: ((), ()),                                          # specimen
         24: ((), ())                                           # isolate
    }

    ZOO_TAX_CODE = {
         1: ((), ("eukaryota")),                                              # domain
         2: ((), ("animalia")),                                               # kingdom
         4: ((), ("chordata", "arthropoda", "mollusca", "nematoda")),         # phylum
         5: ((), ("vertebrata", "myriapoda", "crustacea", "hexapoda")),       # subphylum
         6: ((), ("mammalia", "aves", "reptilia", "amphibia", "insecta")),    # class
         7: ((), ()),                                                         # subclass
         8: ((), ()),                                                         # superorder
         9: ((), ()),                                                         # order
         10: ((), ()),                                                         # suborder
         11: ((), ()),                                                         # infraorder
         12: (("oidea"), ()),                                                 # superfamily
         13: (("oidae"), ()),                                                 # epifamily
         14: (("idae"), ()),                                                  # family
         15: (("inae"), ()),                                                  # subfamily
         16: (("odd"), ()),                                                   # infrafamily
         17: (("ini"), ()),                                                   # tribe
         18: (("ina"), ()),                                                   # subtribe
         19: (("ad", "iti"), ()),                                             # infratribe
         20: ((), ()),                                                        # genus
         21: ((), ()),                                                        # species
         22: ((), ()),                                                        # subspecies
         23: ((), ()),                                                        # specimen
         24: ((), ())                                                         # isolate
    }
    
    VIR_TAX_CODE = {
         2: ((), ("viruses")),      # kingdom
         6: ((), ()),               # class
         7: (("idae"), ()),         # subclass
         9: (("virales"), ()),      # order
         14: (("viridae"), ()),     # family
         15: (("virinae"), ()),     # subfamily
         20: (("virus"), ()),       # genus
         21: ((" virus"), ()),      # species
         23: ((), ()),              # strain
         24: ((), ())               # isolate
    }
    
    TAX_CODE_MAP = {
        "bac": BAC_TAX_CODE,
        "bot": BOT_TAX_CODE,
        "zoo": ZOO_TAX_CODE,
        "vir": VIR_TAX_CODE
    }

    def __init__(self, tax_code_name):
        self.tax_code = TaxCode.TAX_CODE_MAP.get(tax_code_name.lower(), None)
        if not self.tax_code:
            print("ERROR: Unknown taxonomic code: %s" % tax_code_name)
            sys.exit()
        
    @staticmethod    
    def rank_level_name(uni_rank_level):
        return TaxCode.UNI_TAX_RANKS.get(uni_rank_level, ("Unknown", "?__"))
                            
    def guess_rank_level(self, ranks, rank_level):
        rank_name = re.sub("[\W_]+", "", ranks[rank_level].lower())
        
        sorted_tax_levels = sorted(self.tax_code.keys())
        
        real_level = 0

        # first, try to guess rank level based on its name or name suffix
        for lvl in sorted_tax_levels:
          (suffix, exact_match) = self.tax_code[lvl]
          if rank_name.endswith(suffix) or rank_name in exact_match:
            real_level = lvl
            break
            
        # if name-based identification failed, try to derive rank level from its parent
        if real_level == 0:
            if rank_level == 0:    # kingdom
                real_level = 2
            else:
                parent_level = self.guess_rank_level(ranks, rank_level-1)
                if parent_level == 0:
                    return 0
                idx = sorted_tax_levels.index(parent_level)
                for i in range(idx+1, len(sorted_tax_levels)):
                  lvl = sorted_tax_levels[i]
                  suffix = self.tax_code[lvl][0]
                  if lvl in TaxCode.STD_RANKS:
                    real_level = lvl
                    break
                             
        return real_level
         
    def guess_rank_level_name(self, ranks, rank_level):
        real_level = self.guess_rank_level(ranks, rank_level)
        return TaxCode.rank_level_name(real_level)


class Taxonomy:
    EMPTY_RANK = "-"
    RANK_UID_DELIM = "@@"
    
    RANK_PLACEHOLDERS = ["k__", "p__", "c__", "o__", "f__", "g__", "s__"]    

    @staticmethod    
    def lineage_str(ranks):
        return ";".join(ranks).strip(';')

    @staticmethod    
    def lowest_assigned_rank_level(ranks):
        rank_level = len(ranks)-1
        while rank_level >= 0 and ranks[rank_level] == Taxonomy.EMPTY_RANK:
            rank_level -= 1
        return rank_level

    @staticmethod    
    def lowest_assigned_rank(ranks):
        rank_level = Taxonomy.lowest_assigned_rank_level(ranks)
        if rank_level >= 0:
            return ranks[rank_level]
        else:
            return None
        
    @staticmethod    
    def get_rank_uid(ranks, rank_level=-1):
        if rank_level == -1:
            rank_level = Taxonomy.lowest_assigned_rank_level(ranks)        
        return Taxonomy.RANK_UID_DELIM.join(ranks[:rank_level+1])

    @staticmethod    
    def split_rank_uid(rank_uid, min_lvls=0):
        ranks = rank_uid.split(Taxonomy.RANK_UID_DELIM)
        if len(ranks) < min_lvls:
            return ranks + [Taxonomy.EMPTY_RANK] * (min_lvls - len(ranks))
        else:
            return ranks
            
    @staticmethod    
    def rank_uid_to_lineage_str(rank_uid, min_lvls=0):
        ranks = Taxonomy.split_rank_uid(rank_uid, min_lvls)
        return Taxonomy.lineage_str(ranks)

    def __init__(self, prefix="", tax_fname="", tax_map=None):
        self.prefix = prefix
        
        tree_nodes = []
        self.common_ranks = set([])
        self.seq_ranks_map = {}
        self.rank_seqs_map = {}
        self.corr_seq_ids = {}
        self.uncorr_rank_ids = {}
        if tax_map:
            for sid, ranks in tax_map.items():
                self.seq_ranks_map[sid] = ranks
                rank_id = Taxonomy.get_rank_uid(ranks)
                self.rank_seqs_map[rank_id] = self.rank_seqs_map.get(rank_id, []) + [sid]
        elif tax_fname:
            self.load_taxonomy(tax_fname)

    def get_common_ranks(self):
        allranks = list(self.seq_ranks_map.items())
        numitems = len(allranks)
        if numitems > 0:
            self.common_ranks = set(allranks[0][1])
            for i in range(1, numitems):
                curr_set = set(allranks[i][1])
                inters = self.common_ranks & curr_set 
                self.common_ranks = inters
            return self.common_ranks
        else:
            return set([])

    def seq_count(self):
        return len(self.seq_ranks_map)

    def items(self):
        return list(self.seq_ranks_map.items())
        
    def iteritems(self):
        return list(self.seq_ranks_map.items())

    def get_map(self):
        return self.seq_ranks_map

    def remove_seq(self, seqid):
        rank_id = self.seq_rank_id(seqid)
        self.rank_seqs_map[rank_id].remove(seq_id)
        del self.seq_ranks_map[seqid]

    def rename_seq(self, old_seqid, new_seqid):
        rank_id = self.seq_rank_id(old_seqid)
        self.rank_seqs_map[rank_id].remove(old_seqid)
        self.rank_seqs_map[rank_id].append(new_seqid)
        self.seq_ranks_map[new_seqid] = self.seq_ranks_map[old_seqid]
        del self.seq_ranks_map[old_seqid]
        
    def get_seq_ranks(self, seq_id):
        if not seq_id.startswith(self.prefix):
            seq_id = self.prefix + seq_id
        seq_id = self.corr_seq_ids.get(seq_id, seq_id)
        return self.seq_ranks_map[seq_id]
        
    def seq_lineage_str(self, seq_id):
        ranks = list(self.get_seq_ranks(seq_id))
        return Taxonomy.lineage_str(ranks)        

    def seq_rank_id(self, seq_id):
        ranks = list(self.get_seq_ranks(seq_id))
        return Taxonomy.get_rank_uid(ranks)        

    def get_rank_seqs(self, rank_id):
        return self.rank_seqs_map.get(rank_id, [])

    def get_rank_seq_count(self, rank_id):
        return len(self.rank_seqs_map.get(rank_id, []))
        
    def get_uncorr_rank_id(self, rank_id):
        return self.uncorr_rank_ids.get(rank_id, rank_id)
        
    def merge_ranks(self, rank_ids, name_prefix="__TAXCLUSTER__"):
        if len(rank_ids) < 2:
            return None
        
        all_sid = []
        for rank_id in rank_ids:
            all_sid += self.get_rank_seqs(rank_id)
            del self.rank_seqs_map[rank_id]
        
        first_taxon = Taxonomy.split_rank_uid(rank_ids[0])
        merge_lvl = Taxonomy.lowest_assigned_rank_level(first_taxon)
        new_name = name_prefix + first_taxon[merge_lvl]
        new_rank = first_taxon[:merge_lvl] + [new_name]
        new_rank_id = Taxonomy.get_rank_uid(new_rank)

        self.rank_seqs_map[new_rank_id] = all_sid

        for sid in all_sid:
            self.seq_ranks_map[sid] = new_rank
            
        return new_rank_id
        
    def subst_synonyms(self, syn_map):
        for sid, ranks in self.seq_ranks_map.items():
          old_rank_id = Taxonomy.get_rank_uid(ranks)
          changed = False
          for i in range(0, len(ranks)):
            if ranks[i] in syn_map:
                ranks[i] = syn_map[ranks[i]]
                changed = True
                
          if changed:
              new_rank_id = Taxonomy.get_rank_uid(ranks)
#              print old_rank_id, new_rank_id
              del self.rank_seqs_map[old_rank_id]
              self.rank_seqs_map[new_rank_id] = self.rank_seqs_map.get(new_rank_id, []) + [sid]

    def load_taxonomy(self, tax_fname):
        with open(tax_fname) as fin:
            for line in fin:
                line = line.strip()
                toks = line.split("\t")
                sid = self.prefix + toks[0]
                ranks_str = toks[1]
                ranks = ranks_str.split(";")
                for i in range(len(ranks)):
                    rank_name = ranks[i].strip()
    #                if rank_name in GGTaxonomyFile.rank_placeholders:
    #                    rank_name = Taxonomy.EMPTY_RANK
                    ranks[i] = rank_name
                rank_id = Taxonomy.get_rank_uid(ranks)
                    
                self.seq_ranks_map[sid] = ranks 
                self.rank_seqs_map[rank_id] = self.rank_seqs_map.get(rank_id, []) + [sid]

    def normalize_rank_names(self):
        invalid_chars = "[](),;:'"
        sub_chars = "_" * len(invalid_chars)
        trantab = str.maketrans(invalid_chars, sub_chars)
        corr_ranks = {}
        for sid, ranks in self.seq_ranks_map.items():
            old_rank_id = Taxonomy.get_rank_uid(ranks)
            for i in range(len(ranks)):
                if ranks[i] in corr_ranks:
                    ranks[i] = corr_ranks[ranks[i]]
                else:
                    new_rank_name = ranks[i].translate(trantab);
                    if new_rank_name != ranks[i]:
                        corr_ranks[ranks[i]] = new_rank_name
                        ranks[i] = new_rank_name

            new_rank_id = Taxonomy.get_rank_uid(ranks)
            if new_rank_id != old_rank_id and old_rank_id in self.rank_seqs_map:
                self.rank_seqs_map[new_rank_id] = self.rank_seqs_map[old_rank_id]
                del self.rank_seqs_map[old_rank_id]
                self.uncorr_rank_ids[new_rank_id] = old_rank_id
                
        return corr_ranks
        
    def normalize_seq_ids(self):
        invalid_chars = "[](),;:' "
        sub_chars = "_" * len(invalid_chars)
        trantab = str.maketrans(invalid_chars, sub_chars)
        self.corr_seq_ids = {}
        for old_sid in list(self.seq_ranks_map):
            new_sid = old_sid.translate(trantab);
            if new_sid != old_sid:
                self.rename_seq(old_sid, new_sid)
                self.corr_seq_ids[old_sid] = new_sid
                
        return self.corr_seq_ids

    def close_taxonomy_gaps(self):
        for sid, ranks in self.seq_ranks_map.items():
            last_rank = None
            gap_count = 0
            for i in reversed(list(range(1, len(ranks)))):
                if ranks[i] != Taxonomy.EMPTY_RANK:
                    last_rank = ranks[i]
                elif last_rank:
                    gap_count += 1
                    ranks[i] = "parent%d_" % gap_count + last_rank

    def check_for_duplicates(self, autofix=False):
        parent_map = {}
        dups = []
        old_fixed = {}
        old_ranks = {}
        for sid, ranks in self.seq_ranks_map.items():
            i = 1
            while i < len(ranks) and ranks[i] != Taxonomy.EMPTY_RANK:
                parent = ranks[i-1]
                if not ranks[i] in parent_map:
                    parent_map[ranks[i]] = sid
                else:
                    old_sid = parent_map[ranks[i]]
                    if self.get_seq_ranks(old_sid)[i-1] != parent:
                        if autofix:
                            orig_name = self.lineage_str(sid)
                            orig_old_name = self.lineage_str(old_sid)

                            if old_sid not in old_fixed:
                                old_fixed[old_sid] = self.lineage_str(old_sid)
                                old_ranks[old_sid] = set([])
                                
                            if i not in old_ranks[old_sid]:
                                self.seq_ranks_map[old_sid][i] += "_" + self.seq_ranks_map[old_sid][i-1]
                                old_ranks[old_sid].add(i)

                            self.seq_ranks_map[sid][i] = ranks[i] + "_" + parent
                            dup_rec = (old_sid, orig_old_name, sid, orig_name, self.lineage_str(sid))
                        else:
                            dup_rec = (old_sid, self.lineage_str(old_sid), sid, self.lineage_str(sid))
                        dups.append(dup_rec)
                i += 1
                        
        if autofix:
            for sid, orig_name in old_fixed.items():
                dup_rec = (sid, orig_name, sid, orig_name, self.lineage_str(sid))
                dups.append(dup_rec)

        return dups
        
    def check_for_disbalance(self, autofix=False):
        errs = []
        for sid, ranks in self.seq_ranks_map.items():
            if len(ranks) > 7:
                if autofix:
                    orig_name = self.lineage_str(sid)

                    dropq = []
                    keepq = []
                    restq = []
                    for i in range(1, len(ranks)):
                        # drop Subclass and Suborder, preserve Order and Family (based on common suffixes)
                        if ranks[i].endswith("dae") or ranks[i].endswith("neae"):
                            dropq += [i]
                        elif ranks[i].endswith("ceae") or ranks[i].endswith("ales"):
                            keepq += [i]
                        else:
                            restq += [i]

                    to_keep = keepq + restq + dropq
                    to_keep = sorted(to_keep[:-7])
                    
                    new_ranks = [ranks[i] for i in to_keep]
                    self.seq_ranks_map[sid] = new_ranks
                    
                    err_rec = (sid, orig_name, self.lineage_str(sid))
                else:    
                    err_rec = (sid, self.lineage_str(sid))
                errs.append(err_rec)

        return errs
        
        
class TaxTreeBuilder:
    ROOT_LABEL = "<<root>>"

    def __init__(self, config, taxonomy):
        self.tree_nodes = {}
        self.leaf_count = {}
        self.config = config
        self.taxonomy = taxonomy

    def prune_unifu_nodes(self, tree):
        for node in tree.traverse("preorder"):
            if len(node.children) == 1:
                node.delete()

    def add_tree_node(self, tree, nodeId, ranks, rank_level):
        if rank_level >= 0:
            parent_level = rank_level            
            while ranks[parent_level] == Taxonomy.EMPTY_RANK:
                parent_level -= 1
            parentId = Taxonomy.get_rank_uid(ranks, parent_level)
        else:
            parentId = TaxTreeBuilder.ROOT_LABEL
            parent_level = -1

        if (parentId in self.tree_nodes):
            parentNode = self.tree_nodes[parentId]
        else:
            parentNode = self.add_tree_node(tree, parentId, ranks, parent_level-1)
            self.tree_nodes[parentId] = parentNode;

#        print "Adding node: %s, parent: %s, parent_level: %d" % (nodeId, parentId, parent_level)
        newNode = parentNode.add_child()
        newNode.add_feature("name", nodeId)
        return newNode        

    def build(self, min_rank=0, max_seqs_per_leaf=1e9, clades_to_include=[], clades_to_ignore=[]):

        t0 = Tree()
        t0.add_feature("name", TaxTreeBuilder.ROOT_LABEL)
        self.tree_nodes[TaxTreeBuilder.ROOT_LABEL] = t0;
        self.leaf_count[TaxTreeBuilder.ROOT_LABEL] = 0;
        k = 0
        added = 0
        seq_ids = []
        # sequences are leafs of the tree, so they always have the lowest taxonomy level (e.g. "species"+1)        
        for sid, ranks in self.taxonomy.items():
            k += 1
            if self.config.verbose and k % 1000 == 0:
                print("Processed nodes: ", k, ", added: ", added, ", skipped: ", k - added)

            # filter by minimum rank level            
            if ranks[min_rank] == Taxonomy.EMPTY_RANK:
                continue       
    
            # filter by rank contraints (e.g. class Clostridia only)
            clade_is_ok = False

            # check against the inclusion list            
            if len(clades_to_include) > 0:
                for (rank_level, rank_name) in clades_to_include:            
                    if ranks[rank_level] == rank_name:
                        clade_is_ok = True
                        break
            else: # default: include all
                clade_is_ok = True

            # if sequence is about to be included, check it against the ignore list
            if clade_is_ok:
                for (rank_level, rank_name) in clades_to_ignore:
                    if ranks[rank_level] == rank_name:
                        clade_is_ok = False
                        break

            # final decision
            if not clade_is_ok:
                continue

            tax_seq_level = len(ranks)
            parent_level = tax_seq_level - 1
            while ranks[parent_level] == Taxonomy.EMPTY_RANK:
                parent_level -= 1
            parent_name = Taxonomy.get_rank_uid(ranks, parent_level)
            if parent_name in self.tree_nodes:
                parent_node = self.tree_nodes[parent_name]
#                max_seq_per_rank = max_seqs_per_leaf * (tax_seq_level - parent_level)                
                if parent_level == tax_seq_level - 1:
                    max_seq_per_rank = max_seqs_per_leaf # * (tax_seq_level - parent_level)                
                    if parent_name in self.leaf_count and self.leaf_count[parent_name] >= max_seq_per_rank:
                        continue

            self.leaf_count[parent_name] = self.leaf_count.get(parent_name, 0) + 1

            # all checks succeeded: add the sequence to the tree
            self.add_tree_node(t0, sid, ranks, parent_level)
            seq_ids += [sid]
            added += 1

        self.config.log.debug("Total nodes in resulting tree: %d", added)
        
        if self.config.debug:
            reftax_fname = self.config.tmp_fname("%NAME%_mf_unpruned.tre")
            t0.write(outfile=reftax_fname, format=8)

        self.prune_unifu_nodes(t0)
        return t0, seq_ids
