from pyteomics import mztab as mz
import pandas as pd
import re
from utils.search_result_parsing.helper_classes import ProteomicSearchResultReaderBaseClass, InputFile


class MzTabReader(ProteomicSearchResultReaderBaseClass):
    """
        # in column ambiguity_members (PRT) are all protein accessions of one group, the first one of this list is in the
        # 'accession' column of the PRT section, in the PSM section all PSM belonging to protein accessions
        # (column 'accession') are searched
    """
    @staticmethod
    def flatten_list(l):
        return [item for sublist in l for item in sublist]

    @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

    def get_sample_list(self, mtdDict):
        pattern_k = re.compile(r"(ms_run\[(\d+)])(.+)")
        pattern_v = re.compile(r"(ms_run\[(\d+)])")
        samples = set()
        for k in mtdDict.keys():
            try:
                if pattern_k.match(k):
                    samples.add(re.search(pattern_k, k).group(1))
            except TypeError:
                continue
        for v in mtdDict.values():
            try:
                if pattern_v.match(v):
                    re.search(pattern_v, v).group(1)
            except TypeError:
                continue
        sample_list = list(samples)
        sample_list.sort()
        return sample_list

    @staticmethod
    def get_sample_acc_spectra_dict(psmTable):
        """
        :param psmTable: pandas df:  accession | spectra_ref samples | sample
                                 0   BPBMOCPF_02940  ms_run[16]:scan=16536          ms_run[16]
                                1   BPBMOCPF_02940   ms_run[23]:scan=8252          ms_run[23] ...
        :return sample_acc_spectra_dict: dictionary with all PSM IDs per sample and protein accession
        e.g.: {'ms_run[16]': {'BPBMOCPF_02940': {'ms_run[16]:scan=16536'}},
                'ms_run[23]': {'BPBMOCPF_02940': {'ms_run[23]:scan=8252'}, ...
        """
        sample_acc_spectra_dict = {}
        # read all psm into dict
        for i, row in psmTable.iterrows():
            if row['sample'] in sample_acc_spectra_dict.keys():
                if str(row['accession']) in sample_acc_spectra_dict[row['sample']].keys():
                    sample_acc_spectra_dict[row['sample']][str(row['accession'])].add(row['spectra_ref'])
                else:
                    sample_acc_spectra_dict[row['sample']][str(row['accession'])] = {row['spectra_ref']}
            else:
                sample_acc_spectra_dict[row['sample']] = {str(row['accession']): {row['spectra_ref']}}
        return sample_acc_spectra_dict

    def get_spectra_per_protein(self, proteinTable, sample_acc_spectra_dict):
        """
        :param proteinTable:                accession (list all acc pg)  |  ambiguity_members (str all acc per pg)
                            accession(index)
                            BPBMOCPF_03113  ['BPBMOCPF_01459', 'BPBMOCPF_03113', 'JIGCOOEF_0316... | 'BPBMOCPF_03113, BPBMOCPF_01459, JIGCOOEF_03161...'
        :param sample_acc_spectra_dict: e.g. {'ms_run[16]': {'BPBMOCPF_02940': {'ms_run[16]:scan=16536'}},
                                                'ms_run[23]': {'BPBMOCPF_02940': {'ms_run[23]:scan=8252'},..
        :return sample_to_protein_to_spectra_dict: e.g. {'ms_run[16]': {'BPBMOCPF_01459': [''], 'BPBMOCPF_02940':
                                            ['ms_run[16]:scan=16536'], 'JIGCOOEF_03161': [''], 'JIGCOOEF_06185': ['']},
        """
        sample_to_protein_to_spectra_dict = {}
        for sample in sample_acc_spectra_dict.keys():
            sample_to_protein_to_spectra_dict[sample] = {}
            for proteins in proteinTable['accession'].tolist():
                b = False
                for protein in proteins:
                    if protein in sample_acc_spectra_dict[sample].keys():
                        b = True
                        break
                if b:
                    for protein in proteins:
                        try:
                            sample_to_protein_to_spectra_dict[sample][protein] = sorted(sample_acc_spectra_dict[sample][protein])
                        except KeyError:
                            sample_to_protein_to_spectra_dict[sample][protein] = []
        return sample_to_protein_to_spectra_dict

    def get_lists_for_df_building(self, protein_table, sample_list, sample_to_protein_to_spectra_dict):
        quant = []
        spectra_IDs = []
        proteins = []
        samples = []
        for sample in sample_list:
            for accessions in protein_table['accession'].tolist():
                proteins_spectra = []
                for protein in accessions:
                    try:
                        proteins_spectra.append(list(sample_to_protein_to_spectra_dict[sample][protein]))
                    except (KeyError, AttributeError):
                        proteins_spectra.append([])
                if proteins_spectra:
                    spectra_IDs.append(proteins_spectra)
                    proteins.append(accessions)
                    samples.append(sample)
                    quant.append(len(set(self.flatten_list(proteins_spectra))))
        return samples, proteins, quant, spectra_IDs

    def read_file_into_dataframe(self):
        pattern = re.compile(r"(ms_run\[(\d+)]):(.+)")
        # load mztab file
        mztab = mz.MzTab(self.result_table, encoding='utf8', table_format='df')
        meta_data_dict = mztab.metadata
        # from protein section take only accession and ambiguity_members columns,
        proteinTable = mztab.protein_table[['accession', 'ambiguity_members']]
        # from PSM section accession and PSM_ID column
        psmTable = mztab.spectrum_match_table.rename_axis(None).reset_index(drop=True)[['accession', 'spectra_ref']]

        # check ambiguity members
        if not len(set(proteinTable['ambiguity_members'].tolist())) > 1:
            raise AttributeError('Required protein group information in column "ambiguity_members" is not fulfilled.\n')
        # check spectra_ref format
        if not pattern.match(psmTable.iloc[0]['spectra_ref']):
            raise AttributeError('Required spectra_ref format for matching PSMs to sample is not fulfilled.\n' +
                                 'Required format:  ms_run[1-n]:{SPECTRA_REF}\n'+
                                 'Present: ' + psmTable.iloc['spectra_ref'][0])

        sample_list = self.get_sample_list(meta_data_dict)
        #  accession (protein acc) | spectra_ref  (spectra iD) |  sample (ms_run[nb])
        # take ms_run information from spectra_ref into new column sample
        psmTable["sample"] = psmTable['spectra_ref'].apply(lambda x: re.search(pattern, x).group(1))
        # accession (one prot acc) | ambiguity_members (one string all pg accs)
        # ambiguity_members to list of sorted protein_acc united with protein acc from column 'accession'
        proteinTable['accession'] = \
            proteinTable.apply(lambda x: sorted(list(set().union({x['accession']}, x['ambiguity_members']
                                                                 .split(', ')))), axis=1)
        sample_acc_spectra_dict = self.get_sample_acc_spectra_dict(psmTable)
        sample_to_protein_to_spectra_dict = self.get_spectra_per_protein(proteinTable, sample_acc_spectra_dict)
        samples, proteins, quant, spectra_IDs = \
            self.get_lists_for_df_building(proteinTable, sample_list, sample_to_protein_to_spectra_dict)
        # build df from lists
        # sample: ms_run[nb] , proteins: list of proteins in pg, quant: nb spectra per pg per sample,
        # spectra_IDs:  list of list of spectra_IDs per sample/pg
        df_standard = pd.DataFrame({'sample': samples,
                                    'proteins': proteins,
                                    'quant': quant,
                                    'spectra_IDs': spectra_IDs})
        return self.check_and_adapt_accessions(df_standard)
