import sys
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_indexes = []
        for j, elem in enumerate(row_2):
            if elem == 'Master':
                column_indexes.append(j)
            elif elem == 'Accession':
                column_indexes.append(j)
        return column_indexes

    def reduce_full_df_to_abundance_df(self, df, abundance_pos_list):
        abundance_names = [elem[1] for elem in abundance_pos_list]
        column_names = ['Master', 'Accession'] + abundance_names
        columns_indices = self.get_columns_indices(df.iloc[[1]].values.tolist()[0])
        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):
        for i, row in enumerate(df_list[start_row_nb:end_row_nb]):
            if "Master Protein" in row:
                return start_row_nb+i

    def get_mean_abundance(self, protein_group_rows):
        abundances_per_sample = [l[2:] for l in protein_group_rows]
        mean_abundances = [sum(abundance) / len(abundance) for abundance in zip(*abundances_per_sample)]
        return mean_abundances

    def get_protein_group_rows(self, df_list):
        """
        :param df_list: xlsx file read into pandas df as list of lists, one sublist = one row
        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 row[0] == "Master":
                start_row_nb = row_nb + 1
                mst = True
            if type(row[0]) == int and mst:
                end_row_nb = row_nb
                master_rows[pg] = (self.get_master_row_nb(df_list, start_row_nb, end_row_nb))
                pg_to_row_nbs_dict[pg]=(start_row_nb, end_row_nb)
                pg += 1
        # last Protein group
        # last pg not empty
        if type(df_list[-1][0]) != int:
            end_row_nb = len(df_list)
            master_rows[pg] = (self.get_master_row_nb(df_list, start_row_nb, end_row_nb))
            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"]]
        """
        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)
        # get information for protein_groups (row numbers:  start_row_nb - end_row_nb)
        pg_without_informations=0
        pg_without_abundances=[]
        for pg, row_nbs in pg_to_row_nbs_dict.items():
            if row_nbs[0] == row_nbs[1]:
                pg_without_informations += 1
                continue
            samples_per_pg, all_abundance_per_pg = [], []
            # accs protein group
            accs_per_pg = [row[1] 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 master_rows[pg]:
                abundances_per_pg = df_list[master_rows[pg]][2:]
            # no master protein given, mean abundance of all proteins in protein group
            else:
                abundances_per_pg = self.get_mean_abundance(df_list[row_nbs[0]:row_nbs[1]])
            if set(abundances_per_pg) == {0}:
                pg_without_abundances.append(pg)
                continue
            all_abundance_per_pg.extend(abundances_per_pg)
            samples.extend(sample_list)
            abundance.extend(all_abundance_per_pg)
            proteins.extend([accs_per_pg for i in range(len(sample_list))])
        print(f"{pg_without_informations} protein groups without proteins.", file=sys.stderr)
        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)