import sys
from statistics import mean, StatisticsError
import pandas as pd
from utils.search_result_parsing.helper_classes import ProteomicSearchResultReaderBaseClass, InputFile


class PDParser(ProteomicSearchResultReaderBaseClass):
    """
    """
    @staticmethod
    def check_and_adapt_accessions(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

    @staticmethod
    def get_sample_list(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

    @staticmethod
    def get_abundance_columns(df_row_2):
        # Abundance: F12: Sample OR Abundances (Normalized): F9: Sample OR Abundances (Scaled): F7: Sample
        nb_samples = 0
        start_pos_abundance_scales=[]
        start_pos_abundance_normalized=[]
        start_pos_abundance=[]
        for i, elem in enumerate(df_row_2):
            try:
                if elem.startswith('Abundances (Scaled):'):
                    start_pos_abundance_scales.append((i, elem))
                elif elem.startswith('Abundances (Normalized):'):
                    start_pos_abundance_normalized.append((i, elem))
                elif elem.startswith('Abundance:'):
                    start_pos_abundance.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
        else:
            raise ValueError("No abundance information in proteome discoverer group file. Abundance must match one of "
                             "the following: Abundances (Scaled): sample, Abundances (Normalized): sample, Abundance: sample")

    @staticmethod
    def get_columns_indices(row_2):
        column_indexes = []
        for j, elem in enumerate(row_2):
            try:
                if elem == 'Master':
                    column_indexes.append(j)
                elif elem == 'Accession':
                    column_indexes.append(j)
            except AttributeError:
                continue
        return column_indexes

    def read_xlsx_df_into_df(self, df, abundance_pos_list):
        try:
            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])
            df = df.iloc[:, columns_indices]
            df.columns = column_names
            df = df.fillna(0)
            return df
        except FileNotFoundError:
            print('Proteome Discoverer protein group excel file not found.')
            exit(1)

    @staticmethod
    def get_master_row_nb(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_protein_group_rows(self, df_list):
        """
        :param df_list: pd as list: row:
        Master | Accession | Abundances per sample (header of 'sub_df') OR
        Master Protein/Master Protein Candidate/None | accession string | abundances (float) (data of 'sub_df') OR
        'main df': int | int| string ... | 0 ... |
        """
        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
        # not empty last pg
        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, pos_abundance_list):
        """
        :param df: Maser, Accession, Abundance F1, Abundance F2 ...
        :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, 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
        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, 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
            abundance_per_pg.extend(df_list[master_rows[pg]][2:])
            samples.extend(sample_list)
            abundance.extend(abundance_per_pg)
            proteins.extend([accs_per_pg for i in range(len(sample_list))])
        print(f"{pg_without_informations} protein groups without informations.", 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.read_xlsx_df_into_df(df, abundance_pos_list)
        samples, abundance, proteins = self.create_dict_from_xlsx(xlsx_df, sample_list, abundance_pos_list)
        # build df from lists
        df_standard = pd.DataFrame({'sample': samples, 'proteins': proteins, 'quant':abundance})
        return self.check_and_adapt_accessions(df_standard)