from collections import defaultdict
import gzip
from utils.config import INCONCLUSIVE_ANNOTATIONS


class ProcessedTaskData:
    def __init__(self, with_spectra, pgid_to_accs_dict, acc_to_spectra_dict, acc_to_lineages_dict,
                 ignore_unclassified=False, minimum_number_of_annotations=0):
        """
        ## contains all data from task: accessions per pg, spectra per accession and lineage per accession
        ## enables data access (spectra, lineages, accessions, support) by protein group number
        ## check lineage informations and process them (special treatment for 'unclassified' entries)
        :param with_spectra: bool
        :param pgid_to_accs_dict: key: int protein group number, value = list of associated protein accessions
        :param acc_to_spectra_dict: key: string protein accessions value = set of associated spectra IDs
        :param acc_to_lineages_dict: dict key: protein accession value: [[lineage1][lineage2]...]
        :param ignore_unclassified: True -> remove all complete unclassified lineages before analysis
        :param minimum_number_of_annotations: int minimal number of classified entries contained in lineage, remove lineages with less
        # unclassified_number: unique int for distinction of different unclassified
        """
        self.with_spectra = with_spectra
        self.ignore_unclassified = ignore_unclassified
        self.minimum_number_of_annotations = minimum_number_of_annotations
        self.unclassified_number = 0
        self.pgid_to_accs_dict = pgid_to_accs_dict
        self.acc_to_spectra_dict = self._create_acc_to_spectra_dict(acc_to_spectra_dict)
        self.acc_to_lineages_dict = self.check_all_lineages_for_unclassified(acc_to_lineages_dict)

    def check_all_lineages_for_unclassified(self, acc_to_lineages_dict):
        acc_to_clean_lineages_dict = {}
        for acc, lineages in acc_to_lineages_dict.items():
            acc_to_clean_lineages_dict[acc] = self.check_lineages_for_unclassified(lineages)
        return acc_to_clean_lineages_dict

    def _create_acc_to_spectra_dict(self, acc_to_spectra_dict):
        # no spectra in pg_yaml
        if acc_to_spectra_dict is None:
            acc_to_spectra_dict = {}
            for pgid in self.pgid_to_accs_dict.keys():
                for acc in self.pgid_to_accs_dict[pgid]:
                    acc_to_spectra_dict[acc] = {acc}
        return acc_to_spectra_dict

    def _clean_lineages(self, list_of_lineages_per_acc):
        """
        # remove complete lineages: if all entries == 'unclassified' and flag -ignore_unclassified or
        # if less than -minimum_number_of_annotations are classified entries or
        # no information in lineage
        """
        for i, lineage in enumerate(list_of_lineages_per_acc):
            if self.ignore_unclassified:
                if set(lineage) == {'unclassified'}:
                    list_of_lineages_per_acc[i] = []
            if len(set(lineage)) < self.minimum_number_of_annotations+1 and len(set(lineage)) != 1:
                list_of_lineages_per_acc[i] = []
            if set(lineage) in [{''}, {'-'}, {'nan'}]:
                list_of_lineages_per_acc[i] = []
        return list_of_lineages_per_acc

    @staticmethod
    def _replace_undefined_lineage_entries_with_unclassified(lineage):
        return [INCONCLUSIVE_ANNOTATIONS['unclassified'] if elem in ['-', 'nan'] else elem for elem in lineage]

    @staticmethod
    def _add_unique_nb_to_unclassified_entries(lineage, number):
        """
        # if last item in lineage is 'unclassified' -> add unique number to enable distinguishable 'unclassified' entries
        # e.g. 'unclassified' to 'unclassified_1'
        # if not all unclassified support is summed up (even the unclassifieds are different) and get too much support
        """
        return f"{lineage[-1]} {number}"

    @staticmethod
    def _add_last_specific_annotation_to_unclassified_entries(lineage):
        """
        # annotation 'unclassified' added next named annotation, e.g. 'unclassified Ribosomal protein S18', 'unclassified_unclassified_2'
        # to enable a distinction between 'unclassified'  entries
        # E.g.:  ['Alanine metabolism', 'unclassified', 'unclassified' 'K01744'] to
        #        ['Alanine metabolism', 'unclassified K01744', 'unclassified K01744' 'K01744']
        """
        for j, annotation in enumerate(lineage[:-1]):
            if annotation == 'unclassified':
                next_named_annot = lineage[j+1]
                k = j+1
                while next_named_annot == 'unclassified':
                    k += 1
                    if k == len(lineage):
                        break
                    next_named_annot = lineage[k]
                lineage[j] = f"unclassified {next_named_annot}"
        return lineage

    def change_undefined_entries_to_unclassified(self, list_of_lineages_per_acc):
        """
            # replace '-' and 'nan' entries in lineages by 'unclassified'
        """
        for i, lineage in enumerate(list_of_lineages_per_acc):
            if '-' in lineage or 'nan' in lineage:
                list_of_lineages_per_acc[i] = self._replace_undefined_lineage_entries_with_unclassified(lineage)
        return list_of_lineages_per_acc

    def make_unclassified_enties_distinct(self, list_of_lineages_per_acc):
        for i, lineage in enumerate(list_of_lineages_per_acc):
            if 'unclassified' in lineage:
                if lineage[-1] == 'unclassified':
                    list_of_lineages_per_acc[i][-1] = self._add_unique_nb_to_unclassified_entries(lineage, self.unclassified_number)
                    self.unclassified_number += 1
                list_of_lineages_per_acc[i] = self._add_last_specific_annotation_to_unclassified_entries(lineage)
        return list_of_lineages_per_acc

    def check_lineages_for_unclassified(self, list_of_lineages_per_acc):
        """
        :param list_of_lineages_per_acc: all lineages per accession [[][]...]
        :return: list_of_lineages_per_acc, with changed unclassified entries and deleted lineages ([]) depending on params
        minimum_number_of_annotations and ignore_unclassified
        """
        if type(list_of_lineages_per_acc[0]) is not list:
            list_of_lineages_per_acc = [list_of_lineages_per_acc]
        list_of_lineages_per_acc = self._clean_lineages(list_of_lineages_per_acc)
        list_of_lineages_per_acc = self.change_undefined_entries_to_unclassified(list_of_lineages_per_acc)
        list_of_lineages_per_acc = self.make_unclassified_enties_distinct(list_of_lineages_per_acc)
        return list_of_lineages_per_acc

    def get_spectra_set(self, acc):
        return self.acc_to_spectra_dict[acc]

    def get_lineages_list(self, acc):
        return self.acc_to_lineages_dict[acc]

    def get_group_accs(self, pgid):
        return self.pgid_to_accs_dict[pgid]

    def get_group_data(self, pgid):
        accs = self.pgid_to_accs_dict[pgid]
        acc_to_lineages_dict = {}
        accs_to_spectra_dict = {}
        for acc in accs:
            acc_to_lineages_dict[acc] = self.acc_to_lineages_dict[acc]
            accs_to_spectra_dict[acc] = self.acc_to_spectra_dict[acc]
        return accs, acc_to_lineages_dict, accs_to_spectra_dict


class CheckLineageConsistence:
    def __init__(self, path_to_taxmap):
        """
        ## for democratic_lca method inconsistent lineages are possible, to check this and set the inconsistent part to
        ## 'various' the taxmap (taxdump.map.gz) is used. By now only for taxonomic tasks possible
        ##  taxmap file created by  rule make_nr_tax2annot_map.smk
        :param path_to_taxmap:
        """
        self.annot_to_parent_annot_dict = self.load_tax2annot_map(path_to_taxmap)

    @staticmethod
    def load_tax2annot_map(path_to_taxmap):
        annot_to_parent_annot_dict = defaultdict(dict)
        with gzip.open(path_to_taxmap, 'r') as handle:
            handle.readline()
            for line in handle:
                fields = line.decode("utf-8").strip().split('\t')
                lineage = list(reversed(['unclassified' if annot == '-' else annot for annot in fields[2:]]))
                for i, annot in enumerate(lineage[:-1]):
                    annot_to_parent_annot_dict[annot][i] = lineage[i+1]
                annot_to_parent_annot_dict[annot][len(lineage)] = lineage[-1]
        return annot_to_parent_annot_dict

    @staticmethod
    # replaces inconsistent position and all downstream lcas with various
    def replace_with_various(reversed_lineage, inconsistent_pos):
        reversed_lineage[:inconsistent_pos] = ['various']*inconsistent_pos
        return reversed_lineage

    # check consistency of LCA identified lineages, e.g. no eucaryotic species in Bacteria superkingdom
    # load taxmap, start with last lineage item other than unclassified or various
    def check_taxonomic_consistency_of_lineage(self, lineage):
        """
        :param lineage: list of strings e.g. ['Bacteria', ...]
        """
        reversed_lineage = list(reversed(lineage))
        for i, annot in enumerate(reversed_lineage[:-1]):
            if annot in ['unclassified', 'various'] or reversed_lineage[i+1] in ['unclassified', 'various']:
                continue
            try:
                if reversed_lineage[i+1] != self.annot_to_parent_annot_dict[annot][i]:
                    # print('inconsistent lineage %s' % reversed_lineage)
                    reversed_lineage = self.replace_with_various(reversed_lineage, i+1)
                    self.check_taxonomic_consistency_of_lineage(list(reversed(reversed_lineage)))
            # mostly if species name unknown
            except KeyError:
                pass
                # print('taxmap not full %s' % reversed_lineage)
        return list(reversed(reversed_lineage))

