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

# dependencies
from collections import defaultdict
import re
from utils.lca_determination.lca_utils import ProcessedTaskData, CheckLineageConsistence
import pandas as pd

VERSION = '0.0.1'

# TODO if lineages removed by param ignore_unclassified or min_nb_annotation change support: ignore this accs for support determination?
## democratic support: if unclassified not ignored: highest support is possible for unclassified entry,
## democratic: second support value from support group_support or democratic_annot_weight dict?
class GroupLCA:
    def __init__(self, pgid, nb_levels, all_group_data, method, is_tax_task, democratic_annotation_weight_df=None, threshold=None, checker=None):
        """
        ## for democratic_lca method first a annotation_weight_df for all data must be generated
        ## for group and democratic lca method annotaion_weight dict are generated per group
        ## all data per task (lineages, groups, accessions, spectra) in RAM, for democratic_lca also democratic_annotation_weight_df
        :param pgid: int protein_group ID
        :param nb_levels: number levels
        :param all_group_data = acc_to_lineages_dict: keys = accession, values [[lineage1][lineage2]..] per group,
            pgid_to_accs_dict: key: int protein group ID, value = list of protein accessions per group,
            accs_to_spectra_dict: key: str protein accession, value= set of spectra IDs per group
        :param method: LCA method, 'lca' or 'democratic_lca'
        :param is_tax_task = True if taxonomic task
        :param democratic_annotation_weight_df: pre calculated democratic_annotation_weight_df if method = 'democratic_lca'
        :param threshold: threshold for group LCA
        :param checker: for tax task, checking lineage consistancy
        """
        self.nb_levels = nb_levels
        self.method = method
        self.threshold = threshold
        self.is_tax_task = is_tax_task
        self.pgid = pgid
        self.acc_to_lineages_dict = all_group_data[1]
        self.accs = all_group_data[0]
        self.accs_to_spectra_dict = all_group_data[2]
        self.group_support = len(set(self._flatten_list(self.accs_to_spectra_dict.values())))
        self.checker = checker
        if method == 'lca':
            self.annotation_weight_dict = self.get_annot_weight_dict_per_group()
        else:
            self.annotation_weight_df = democratic_annotation_weight_df
            self.annotation_weight_dict_group = self.get_annot_weight_dict_per_group()

    @staticmethod
    def _clean_name(name):
        """
        :param name: level_nb_annotation
        :return: clean name without number
        """
        return re.sub("^[0-9]+ ", "", name)

    @staticmethod
    def _flatten_list(list_of_lists):
        return [val for sublist in list_of_lists for val in sublist]

    @staticmethod
    def get_max_indices(support_per_annot):
        return [index for index, value in enumerate(support_per_annot) if value == max(support_per_annot)]

    @staticmethod
    def _get_part_dict(dict, key_list):
        return {key: dict[key] for key in key_list}

    @staticmethod
    def is_all_unclassified(annot_list):
        return {'unclassified' if 'unclassified' in name else name for name in annot_list} == {'unclassified'}

    def _find_max_value(self, level_nb_annot_to_support_dict):
        """
        :param level_nb_annot_to_support_dict: dict of key : numbered annot, value=number of assigned spectra
                    e.g.{'0 Bacteria': 4, '0 Eukaryota': 3}
        :return list of nb_annot(/s) with max_value
        """
        k = list(level_nb_annot_to_support_dict.keys())
        indices = self.get_max_indices(level_nb_annot_to_support_dict.values())
        return [k[i] for i in indices]

    @staticmethod
    def get_numbered_annotation_lineage(lineage):
        level_numbered_annotations = set()
        for i, annot in enumerate(lineage):
            level_numbered_annotations.add(f"{i} {annot}")
        return level_numbered_annotations

    @staticmethod
    def _get_annot_dict_by_level_nb(level_nb, annotation_dict):
        return dict((k, v) for k, v in annotation_dict.items() if k.startswith(str(level_nb)))

    def _get_annot_list_by_level_nb(self, level_nb):
        return [level_name for level_name in self.annotation_weight_dict.keys() if
                level_name.startswith(str(level_nb))]

    def get_level_number_annotations_per_acc(self, acc):
        """
        :param acc: accession
        :return level_numbered_annotations, set of annotations extended with level number e.g. {'1 Bacteria', ...}
        """
        level_numbered_annotations = set()
        for lineage in self.acc_to_lineages_dict[acc]:
            level_numbered_annotations = level_numbered_annotations.union(self.get_numbered_annotation_lineage(lineage))
        return level_numbered_annotations

    def get_level_number_annotations_per_group(self):
        level_numbered_annot = set()
        for acc in self.accs:
            level_numbered_annot = level_numbered_annot.union(self.get_level_number_annotations_per_acc(acc))
        return level_numbered_annot

    def get_annot_weight_dict_per_group(self):
        """
        :return annot_to_support_dict: key = 'level_nb annotation' (e.g. '6 Vibrio sp. 2017V-1176')
                    value= int number of (protein or spectral IDs) supporting annot
        """
        annot_to_spectra_dict = defaultdict(set)
        for acc in self.accs:
            annot_set = self.get_level_number_annotations_per_acc(acc)
            for annot in annot_set:
                annot_to_spectra_dict[annot] = annot_to_spectra_dict[annot].union(self.accs_to_spectra_dict[acc])
        annot_to_support_dict = {key: len(value) for key, value in annot_to_spectra_dict.items()}
        return annot_to_support_dict

    def get_democratic_support(self, nb_annotation):
        return self.annotation_weight_df.loc[nb_annotation, 'spectra']

    def get_annot_support(self, lca):
        return self.annotation_weight_dict_group[lca]

    def get_max_annot_support(self, max_list):
        indices = self.get_max_indices([self.annotation_weight_dict_group[max_lca] for max_lca in max_list])
        return self.get_democratic_support(max_list[indices[0]])

    def set_lca_and_support_lineage_to_unclassified(self):
        level_nb_lca_set = {f"{level_nb} unclassified" for level_nb in range(self.nb_levels)}
        lca_to_support_dict = dict.fromkeys(level_nb_lca_set, '0/0')
        return level_nb_lca_set, lca_to_support_dict

    @staticmethod
    def set_lca_and_support_to_unclassified(level_nb_lca_set, lca_to_support_dict, level_nb):
        lca = f"{level_nb} unclassified"
        level_nb_lca_set.add(lca)
        lca_to_support_dict[lca] = '0/0'
        return level_nb_lca_set, lca_to_support_dict

    ################  group lca functions  #######################################################################

    def _get_annotations_above_threshold_to_support_dict(self):
        return dict((k, v) for k, v in self.annotation_weight_dict.items() if
                    v >= round((self.group_support * self.threshold), 4))

    def _set_lca_to_unclassified_or_various(self, level_nb, nb_annot_above_th_to_support_dict=None):
        level_names = self._get_annot_list_by_level_nb(level_nb)
        if nb_annot_above_th_to_support_dict:
            annot_weight_dict = nb_annot_above_th_to_support_dict
        else:
            # get member highest support
            annot_weight_dict = self._get_part_dict(self.annotation_weight_dict, level_names)

        max_list = self._find_max_value(annot_weight_dict)
        if self.is_all_unclassified(level_names):
            return max_list[0], f"{level_nb} unclassified"
        else:
            return max_list[0], f"{level_nb} various"

    def get_group_lca_and_support(self, return_highest=True):
        """
        :param return_highest: return only highest matches (of different above threshold)
        :return: level_nb_lca_set  e.g.: {'0 Bacteria', `1 Firmicutes', 2 unclassified, 3 various ...}
        :return: lca_to_support_dict e.g.:{'0 Bacteria': '10/12', 2 unclassified: '(2)/12'...}
        """

        all_accessions = self._flatten_list(self._flatten_list((self.acc_to_lineages_dict.values())))
        if self.is_all_unclassified(all_accessions):
            return self.set_lca_and_support_lineage_to_unclassified()

        group_annotations_above_th_dict = self._get_annotations_above_threshold_to_support_dict()
        level_nb_lca_set = set()
        lca_to_support_dict = {}

        for level_nb in range(self.nb_levels):
            # group with all unclassified -> empty list_of_lineages_per_acc after check_lineage_for_unclassified
            if not self.annotation_weight_dict:
                level_nb_lca_set, lca_to_support_dict = self.set_lca_and_support_to_unclassified(level_nb_lca_set, lca_to_support_dict, level_nb)
                continue

            nb_annot_above_th_to_support_dict = self._get_annot_dict_by_level_nb(level_nb, group_annotations_above_th_dict)
            # only one with value about threshold
            if len(nb_annot_above_th_to_support_dict) == 1:
                lca = list(nb_annot_above_th_to_support_dict.keys())[0]
                lca_to_support_dict[lca] = f"{self.annotation_weight_dict[lca]}/{self.group_support}"
                level_nb_lca_set.add(lca)
            # no entry for level with value > threshold
            elif not nb_annot_above_th_to_support_dict:
                max_lca, lca = self._set_lca_to_unclassified_or_various(level_nb)
                level_nb_lca_set.add(lca)
                lca_to_support_dict[lca] = f"({self.annotation_weight_dict[max_lca]})/{self.group_support}"
            # multiple entries with value > threshold
            else:
                max_lca, lca = self._set_lca_to_unclassified_or_various(level_nb, nb_annot_above_th_to_support_dict)
                if return_highest:
                    level_nb_lca_set.add(lca)
                    lca_to_support_dict[lca] = f"{self.annotation_weight_dict[max_lca]}/{self.group_support}"
                # if multiple entries with value > threshold return various, spectra-list
                else:
                    lca = f"{level_nb} various"
                    lca_to_support_dict[lca] = f"{self.annotation_weight_dict[max_lca]}/{self.group_support}"
                    level_nb_lca_set.add(lca)

        return level_nb_lca_set, lca_to_support_dict

    ################  democratic lca functions  #######################################################################
    @staticmethod
    def get_annot_per_level(level_numbered_annot, level_nb):
        return [annot for annot in level_numbered_annot if annot.startswith(str(level_nb))]

    def _get_lcas_with_max_support(self, level_numbered_annot, level_nb):
        """

        """
        annot_per_level_list = self.get_annot_per_level(level_numbered_annot, level_nb)
        democratic_support_per_annot = list(self.annotation_weight_df.loc[annot_per_level_list, 'spectra'])
        indices_of_annot_with_max_support = self.get_max_indices(democratic_support_per_annot)
        return [annot_per_level_list[i] for i in indices_of_annot_with_max_support]

    # TODO: insert threshold?
    def get_democratic_lca_and_support(self):
        """
        :param pgid: int protein group ID
        :return level_to_lca_with_spectra_set_dict: key: level, value: [('annotation1', int}), 'annotation2..]
         with highest support in all data  e.g. {'superkingdom': [('Bacteria', 2)],
                                                    'phylum': [('Firmicutes', 5)],...]}
        :return: lca_to_support_dict e.g.:{'0 Bacteria': '10/200', 2 unclassified: '1/12'...}
        """
        all_accessions = self._flatten_list(self._flatten_list((self.acc_to_lineages_dict.values())))
        if self.is_all_unclassified(all_accessions):
            return self.set_lca_and_support_lineage_to_unclassified()

        level_numbered_annot = self.get_level_number_annotations_per_group()
        level_nb_lca_set = set()
        lca_to_support_dict = {}
        for level_nb in range(self.nb_levels):
            if not self.annotation_weight_dict_group or level_numbered_annot == set():
                level_nb_lca_set, lca_to_support_dict = self.set_lca_and_support_to_unclassified(level_nb_lca_set, lca_to_support_dict, level_nb)
            else:
                max_lca_list_per_level = self._get_lcas_with_max_support(level_numbered_annot, level_nb)
                if len(max_lca_list_per_level) > 1:
                    if self.is_all_unclassified(max_lca_list_per_level):
                        lca = f"{level_nb} unclassified"
                        level_nb_lca_set.add(lca)
                        lca_to_support_dict[lca] = f"{self.get_max_annot_support(max_lca_list_per_level)}/" \
                                                   f"{self.get_democratic_support(max_lca_list_per_level[0])}"
                    else:
                        lca = f"{level_nb} various"
                        level_nb_lca_set.add(lca)
                        lca_to_support_dict[lca] = f"{self.get_max_annot_support(max_lca_list_per_level)}/" \
                                                   f"{self.get_democratic_support(max_lca_list_per_level[0])}"
                else:
                    lca = max_lca_list_per_level[0]
                    level_nb_lca_set.add(lca)
                    lca_to_support_dict[lca] = f"{self.get_annot_support(lca)}/" \
                                               f"{self.get_democratic_support(lca)}"
        return level_nb_lca_set, lca_to_support_dict

   ####################################################################################################################
    def check_lineage_consistency(self, lca_lineage):
        lca_lineage = self.checker.check_taxonomic_consistency_of_lineage(lca_lineage)
        lca_list = list(self.get_numbered_annotation_lineage(lca_lineage))
        lca_list.sort()
        return [self._clean_name(lca) for lca in lca_list]

    def get_lca_and_support(self):
        if self.method == 'lca':
            lca_set, lca_to_support_dict = self.get_group_lca_and_support()
        elif self.method == 'democratic_lca':
            lca_set, lca_to_support_dict = self.get_democratic_lca_and_support()
        else:
            raise ValueError(f"Invalid lca method '{self.method}' defined. Allowed methods: 'lca' and 'democratic_lca'")

        lca_list = list(lca_set)
        lca_list.sort()
        lca_lineage = [self._clean_name(lca) for lca in lca_list]
        if self.is_tax_task and self.checker is not None:
            lca_lineage = self.check_lineage_consistency(lca_lineage)

        support_list = []
        for lca in lca_list:
            support_list.append(lca_to_support_dict[lca])

        lca_lineage = ['unclassified' if re.match("unclassified .*", i) else i for i in lca_lineage]
        return lca_lineage, support_list


class WeightedLCA:
    def __init__(self, lca_level_names, acc_to_lineages_dict, pgid_to_accs_dict, accs_to_spectra_dict, with_spectra,
                 method, path_to_taxmap, is_tax_task, threshold=1, ignore_unclassified=False,
                 minimum_number_of_annotations=0):
        """
        ## for democratic_lca method first a annotation_weight_df for all data must be generated (RAM!)
        ## for group lca method annotaion_weight dict are generated per group
        ## all data per task (lineages, groups, accessions, spectra) in RAM
        :param lca_level_names: list of level names, e.g. ['superkingdom', 'family',...]
        :param acc_to_lineages_dict: keys = accession, values [[lineage1][lineage2]..]
        :param pgid_to_accs_dict: key: int protein group ID, value = list of protein accessions
        :param accs_to_spectra_dict: key: str protein accession, value= set of spectra IDs
        :param with_spectra: True if spectra IDs contained in protein_group.yaml
        :param method: LCA method, 'lca' or 'democratic_lca'
        :param path_to_taxmap: Path to Taxmap for checking lineage consistency for tax tasks
        :param threshold: threshold for group LCA, democratic LCA need no threshold
        :param ignore_unclassified: if true, all complete unclassified lineages are ignored
        :param minimum_number_of_annotations: minimum number of annotations (not 'unclassified') for including lineage in lca determination
        e.g. =2, only lineages with at least 2 classified are considered, else removed for analysis
        ## all_groups_data: TaskDataPerProteinGroup object (utils)
        """
        self.lca_level_annotations = lca_level_names
        self.method = method
        self.threshold = threshold
        self.is_tax_task = is_tax_task
        self.all_groups_data = ProcessedTaskData(with_spectra, pgid_to_accs_dict, accs_to_spectra_dict,
                                                 acc_to_lineages_dict, ignore_unclassified, minimum_number_of_annotations)
        self.total_number_of_pgs = len(self.all_groups_data.pgid_to_accs_dict.keys())
        self.path_to_taxmap = path_to_taxmap
        self.democratic_annotation_weight_df = self.get_democratic_annot_weight_df(self.all_groups_data) \
            if method == 'democratic_lca' else None

    @staticmethod
    def write_headers(output_lca, output_lca_support, lca_headline):
        """
        :param output_lca: path to lca output
        :param output_lca_support: path to lca_support output
        :param lca_headline headline from task
        """
        with open(output_lca, 'w') as support_handle:
            support_handle.write(f"#pg\t{lca_headline}\n")
        with open(output_lca_support, 'w') as support_handle:
            support_handle.write(f"#pg\t{lca_headline}\n")

    @staticmethod
    def write_part_lca_result(pgid_to_lcas_dict, output_lca):
        """
        # write lca file
        :param pgid_to_lcas_dict
        :param output_lca: path to lca output
        """
        with open(output_lca, 'a') as handle:
            for pgid in pgid_to_lcas_dict.keys():
                handle.write(f"{pgid}\t" + '\t'.join(pgid_to_lcas_dict[pgid]) + "\n")

    @staticmethod
    def _add_level_number_to_annotations(lineages):
        """
        :param lineages: list of list [[lineage1][lineage2]]
        :return level_numbered_annotations, set of annotations extended with level number e.g. {'1 Bacteria', ...}
        """
        level_numbered_annotations = set()
        for lineage in lineages:
            for i, annot in enumerate(lineage):
                level_numbered_annotations.add(f"{i} {annot}")
        return level_numbered_annotations

    def get_democratic_annot_weight_df(self, all_group_data):
        """
        create dataframe with level numbered annotations (e.g. '0_Bacteria') und column spectra containing set of
        spectra as concatenated string
        :param all_group_data = self.all_groups_data
        :return: annot_weight_df
        """
        nb_annot_list = []
        spectra_list = []
        for pgid in range(self.total_number_of_pgs):
            for acc in all_group_data.get_group_accs(pgid):
                lineages = all_group_data.get_lineages_list(acc)
                spectra = all_group_data.get_spectra_set(acc)
                nb_annot = self._add_level_number_to_annotations(lineages)
                nb_annot_list.extend(list(nb_annot))
                for i in range(len(nb_annot)):
                    spectra_list.append('; '.join(spectra))
        annot_weight_df = pd.DataFrame(list(zip(nb_annot_list, spectra_list)), columns=['nb_annot', 'spectra'])
        annot_weight_df = annot_weight_df.groupby(annot_weight_df['nb_annot'])['spectra'].apply(lambda x: '; '.join(x)).reset_index()
        annot_weight_df['spectra'] = annot_weight_df['spectra'].apply(lambda x: len(set(x.split('; '))))
        annot_weight_df = annot_weight_df.set_index('nb_annot')
        return annot_weight_df

    def weight_and_write_lca(self, output_lca, output_lca_support):
        """
        determines lca and writes lca, lca-support file for all protein groups in 1000er blocks
        :param output_lca: path to lca output
        :param output_lca_support: path to lca_support output
        """
        pgid_to_lcas_dict = {}
        pgid_to_supports_dict = {}
        lca_headline = '\t'.join(self.lca_level_annotations)
        self.write_headers(output_lca, output_lca_support, lca_headline)
        checker = CheckLineageConsistence(self.path_to_taxmap) if self.is_tax_task and self.path_to_taxmap is not None else None

        for pgid in range(self.total_number_of_pgs):
            lca_per_group = GroupLCA(pgid, len(self.lca_level_annotations), self.all_groups_data.get_group_data(pgid),
                                     self.method, self.is_tax_task, self.democratic_annotation_weight_df, self.threshold, checker=checker)
            result = lca_per_group.get_lca_and_support()
            pgid_to_lcas_dict[pgid] = result[0]
            pgid_to_supports_dict[pgid] = result[1]

            # write 1000 lca results
            if pgid % 1000 == 0 and pgid != 0:
                self.write_part_lca_result(pgid_to_lcas_dict, output_lca)
                self.write_part_lca_result(pgid_to_supports_dict, output_lca_support)
                pgid_to_lcas_dict = {}
                pgid_to_supports_dict = {}
        self.write_part_lca_result(pgid_to_lcas_dict, output_lca)
        self.write_part_lca_result(pgid_to_supports_dict, output_lca_support)
