import sys
from collections import OrderedDict

import pandas as pd
from utils.search_result_parsing.helper_classes import ProteomicSearchResultReaderBaseClass, InputFile


class PDParser(ProteomicSearchResultReaderBaseClass):
    """
    parsing of the xlsx protein group proteome discoverer output:
    1. read excel file into pandas dataframe
    dataframe structure:
        header: used for counting samples (header elements starting with "Found in Sample" are counted)
        protein_group_header: with elements "Master", "Accession" and "Abundance...: sample "
        protein_group_rows: all rows until next protein_group_header. Contain following information:
         [acc, master-protein/master-candidate/None, abundance values per sample]
    2. determine samples -> counting of "Found in Sample" in df header -> sample names = "F+count_nb+1"
    3. determine columns used for abundance: first protein group header (=row 2) must match one of this terms
            ['Abundance:', 'Abundances (Normalized):', 'Abundances (Scaled):']
            If multiple matching Abundance terms are present, only one is selected, in the order above
    4. determine column position of needed columns (Master, Accession, Abundances) and reduce df to these columns
    5. determine group rows for every protein group:
            determine position of first protein group row and last protein group row in df
    6. get quant values: if master-protein is present (mostly): abundances taken from the master-protein row of a protein group,
        if no abundance entries for all samples in the master-protein row --> protein group discarded
        if master_protein not present: abundances = mean(abundances of all proteins of this group)
    """
    def check_and_adapt_accessions(self, df_standard):
        for accessions in df_standard.proteins:
            for i, accession in enumerate(accessions):
                if len(accession.split('|')) > 2:
                    accessions[i] = accession.split('|')[1]
        return df_standard

    def get_sample_list(self, df_header):
        nb_samples = 0
        for elem in df_header:
            if "Found in Sample" in elem:
                nb_samples += 1
        samples = [f"F{nb+1}" for nb in range(nb_samples)]
        return samples

    def get_abundance_columns(self, df_row_2):
        """
        Find positions of Abundances of different samples, if multiple different abundance methods existis, following
        abundances are used in this order: Abundances, Abundances (Scaled), Abundances (Normalized)
        :param df_row_2: second row of dataframe
        :return start_pos_abundance: list of tuples, tuples contain the index (=column number) of the abundance values
         and the column content, e.g. [(12, "Abundances F1: Sample"), (13, "Abundances F2: Sample"), ...]
        """
        start_pos_abundance_scales=[]
        start_pos_abundance_normalized=[]
        start_pos_abundance=[]
        for i, elem in enumerate(df_row_2):
            try:
                if elem.startswith('Abundance:'):
                    start_pos_abundance.append((i, elem))
                elif elem.startswith('Abundances (Normalized):'):
                    start_pos_abundance_normalized.append((i, elem))
                elif elem.startswith('Abundances (Scaled):'):
                    start_pos_abundance_scales.append((i, elem))
            except AttributeError:
                continue
        if len(start_pos_abundance) > 0:
            return start_pos_abundance
        elif len(start_pos_abundance_scales) > 0:
            return start_pos_abundance_scales
        elif len(start_pos_abundance_normalized) > 0:
            return start_pos_abundance_normalized
        else:
            raise ValueError("No abundance information in proteome discoverer group file. Abundance must match one of "
                             "the following:  Abundance: sample, Abundances (Scaled): sample, Abundances (Normalized): sample")

    def get_columns_indices(self, row_2):
        """
        param row_2: second row = first header_row of "PG sub_df"
        """
        column_to_index_dict = OrderedDict()
        for j, elem in enumerate(row_2):
            if elem == 'Master':
                column_to_index_dict['Master'] = j
            elif elem == 'Accession':
                column_to_index_dict['Accession'] = j
        if "Accession" not in column_to_index_dict.keys():
            raise AssertionError('Accession column is missing in protein groups!')
        if 'Master' not in column_to_index_dict.keys():
            print('Master column is not available. Instead of using abundance values of the master protein,  '
                  'non-zero abundances of all protein group members are averaged.', file=sys.stderr)
        return column_to_index_dict

    def reduce_full_df_to_abundance_df(self, df, abundance_pos_list):
        column_to_index_dict = self.get_columns_indices(df.iloc[[1]].values.tolist()[0])
        column_names = [k for k in column_to_index_dict.keys()]
        columns_indices = [v for v in column_to_index_dict.values()]
        abundance_names = [elem[1] for elem in abundance_pos_list]
        column_names.extend(abundance_names)
        columns_indices.extend([elem[0] for elem in abundance_pos_list])
        abundance_df = df.iloc[:, columns_indices]
        abundance_df.columns = column_names
        abundance_df = abundance_df.fillna(0)
        return abundance_df

    def get_master_row_nb(self, df_list, start_row_nb, end_row_nb, pg):
        for i, row in enumerate(df_list[start_row_nb:end_row_nb]):
            if "Master Protein" in row:
                return start_row_nb+i
        raise KeyError(f'No "Master Protein" row in protein group {pg} (row {start_row_nb} -  row {end_row_nb})')

    def get_mean_abundance(self, protein_group_rows, abundance_start_index):
        abundances_per_sample = [l[abundance_start_index:] for l in protein_group_rows]
        mean_abundances = []
        for abundance in zip(*abundances_per_sample):
            # remove protins with no accession from mean calculation
            abundance_without_zero = [a for a in abundance if a !=0]
            if len(abundance_without_zero)==0:
                mean_abundances.append(0)
            else:
                mean_abundances.append(sum(abundance_without_zero)/len(abundance_without_zero))
        return mean_abundances

    def get_protein_group_rows(self, df_list, with_master_column):
        """
        :param df_list: xlsx file read into pandas df as list of lists, one sublist = one row
        :param with_master_column: bool master column in data
        the xlsx file contains 3 different types of rows:
        1. "main_df row": contains no nedded protein group information-> rows were only used as mark for end of a protein goup
        2. header "PG sub_df":  [Master, Accession, Abundances per sample ](header of 'PG sub_df'), used als mark for start of a new protein group
        3. rows "PG sub_df": contains the relevant abundances (float) and protein accession information
        :return pg_to_row_nbs_dict: key= pg, value=(list index of the first PG row of one PG, list index of the last PG row of one PG)
        :return master_rows: dict key=pg number, value=list index of the PG row, where value of Master column ("PG sub_df")=="Master Protein"
        """
        pg_to_row_nbs_dict = {}
        master_rows = {}
        mst = False
        pg = 0
        # find end and start row per protein group
        for row_nb, row in enumerate(df_list):
            if "Accession" in row: # new group header
                start_row_nb = row_nb + 1
                mst = True
            if type(row[0]) == int and mst:
                end_row_nb = row_nb
                pg_to_row_nbs_dict[pg]=(start_row_nb, end_row_nb)
                if with_master_column:
                    try:
                        master_rows[pg] = (self.get_master_row_nb(df_list, start_row_nb, end_row_nb, pg))
                    # no Master row
                    except KeyError:
                        master_rows[pg] = None
                pg += 1
        # last Protein group
        # last pg not empty
        if type(df_list[-1][0]) != int:
            end_row_nb = len(df_list)
            if with_master_column:
                try:
                    master_rows[pg] = (self.get_master_row_nb(df_list, start_row_nb, end_row_nb, pg))
                except KeyError:
                    master_rows[pg] = None
            pg_to_row_nbs_dict[pg] = (start_row_nb, end_row_nb)
        return pg_to_row_nbs_dict, master_rows

    def create_dict_from_xlsx(self, df, sample_list):
        """
        :param df: dataframe with columns: (Master), Accession, Abundance sample 1, Abundance sample 2 ...
        :param sample_list: List of sorted sample names ["F1", "F10", "F11"..]
        : return samples, abundance, proteins: lists, columns in df_standard
        # e.g. ["F1", "F2"], [100, 45.5], [["acc1", "acc2"], ["acc1", "acc2"]]
        """
        with_master_column = 'Master' in df.columns
        accession_column_index = list(df).index('Accession')
        abundance_start_index = list(df).index(next(x for x in list(df) if 'Abundance' in x))
        df_list = df.values.tolist()
        samples, proteins, abundance = [], [], []
        # pg_to_row_nbs_dict: dict{pg: ()}; master_rows: dict{pg: line number} where column Master == 'Master Protein'
        pg_to_row_nbs_dict, master_rows = self.get_protein_group_rows(df_list, with_master_column)
        # get information for protein_groups (row numbers:  start_row_nb - end_row_nb)
        pg_without_information=0
        pg_without_abundances=[]
        for pg, row_nbs in pg_to_row_nbs_dict.items():
            if row_nbs[0] == row_nbs[1]:
                pg_without_information += 1
                continue
            samples_per_pg, all_abundances_per_pg = [], []
            # accs protein group
            accs_per_pg = [row[accession_column_index] for row in df_list[row_nbs[0]:row_nbs[1]]]
            # master protein to first position (not possible, sorted by alphabet)
            # accs_per_pg.insert(0, accs_per_pg.pop(accs_per_pg.index(pg_master)))
            # abundance: row abundance values of master protein per protein group per sample

            # take abundance only for master protein
            if with_master_column and master_rows[pg]:
                all_abundances_per_pg = df_list[master_rows[pg]][abundance_start_index:]
            # no master protein given, mean abundance of all proteins in protein group
            else:
                all_abundances_per_pg = self.get_mean_abundance(df_list[row_nbs[0]:row_nbs[1]], abundance_start_index)
            if set(all_abundances_per_pg) == {0}:
                pg_without_abundances.append(pg)
                continue
            samples.extend(sample_list)
            abundance.extend(all_abundances_per_pg)
            proteins.extend([accs_per_pg for i in range(len(sample_list))])
        print(f"{pg_without_information} protein groups without proteins.", file=sys.stderr)
        if pg_without_abundances:
            print(f"Following protein groups are without abundance informations: {', '.join([str(pg) for pg in (pg_without_abundances)])}", file=sys.stderr)
        return samples, abundance, proteins

    def read_file_into_dataframe(self):
        df = pd.read_excel(self.result_table, engine="openpyxl")
        df = df.dropna(how='all')
        sample_list = self.get_sample_list(df.columns)
        abundance_pos_list = self.get_abundance_columns(df.iloc[1].tolist())
        xlsx_df = self.reduce_full_df_to_abundance_df(df, abundance_pos_list)
        samples, abundance, proteins = self.create_dict_from_xlsx(xlsx_df, sample_list)

        # build df from lists
        df_standard = pd.DataFrame({'sample': samples, 'proteins': proteins, 'quant': abundance})
        return self.check_and_adapt_accessions(df_standard)