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

import pandas as pd
from pyteomics import mzid # This module requires lxml
from utils.search_result_parsing.helper_classes import ProteomicSearchResultReaderBaseClass


class MzidentTableReader(ProteomicSearchResultReaderBaseClass):
    """
    parse mzIdentML 

    minimum required mzIdentML version: 1.2, format specification: https://www.psidev.info/mzidentml
    
    Required mzIdentML tags:
    - ProteinAmbiguityGroup
    - ProteinDetectionHypothesis
    - PeptideDetectionHypothesis
    - SpectrumIdentificationResults
    
    Data are processed into protein groups, spectra and quant as follows:
    - protein_group: created based on all accs (acc of corresponding dBSequence_ref) from one ProteinAmbiguityGroup
    - spectra: IDs of all SpectrumIdentificationResults of one ProteinDetectionHypothesis in one list
        (IDs retrieved from ProteinAmbiguityGroup -> ProteinDetectionHypothesis -> PeptideDetectionHypothesis
         -> SpectrumIdentificationItem -> SpectrumIdentificationResult(=xml parent), if passThreshold=True or not contained for all levels)
    - quant: number of different spectra IDs of all ProteinDetectionHypothesis of one ProteinAmbiguityGroup
    """

    def check_style(self):
        pass

    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_spectra_id_SpectrumIdentificationItem_dict(self, reader):
        sii_to_spectra_id_and_ref = {}
        for ident_result in reader.iterfind("SpectrumIdentificationResult"):
            for sii in ident_result['SpectrumIdentificationItem']:
                sii_to_spectra_id_and_ref[sii['id']] = (ident_result['spectrumID'], ident_result['spectraData_ref'])
        return sii_to_spectra_id_and_ref

    def create_pg_df(self, reader, sii_to_spectra_id_and_ref):
        """
        :param reader: pyteomics mzident readerbased on index
        :param sii_to_spectra_id_and_ref: dict {spectrum_identification_item: (spectra_ID, spectra_ref)}
        :return df_pg: df with columns sample | protein_group | accession | spectrum_identification
        """
        pg_list = []
        for pag in reader.iterfind("ProteinAmbiguityGroup"):
            # cvParam MS:1002415 =  “protein group passes threshold” MUST be contained (but is not always)
            if ('protein group passes threshold' in pag) and (pag['protein group passes threshold'] is False):
                continue
            pg_id = pag['id']
            for pdh in pag['ProteinDetectionHypothesis']:
                # value passThreshold MAY be contained
                if ('passThreshold' in pdh) and (pdh['passThreshold'] is False):
                    continue
                seq_accession = reader.get_by_id(pdh['dBSequence_ref']).get("accession")
                # Map peptide evidence to the spectra that they are assigned to
                for pep in pdh['PeptideHypothesis']:
                    for ref in pep['SpectrumIdentificationItemRef']:
                        sii = reader.get_by_id(ref['spectrumIdentificationItem_ref'])
                        if ('passThreshold' in sii) and (sii['passThreshold'] is False):
                            continue
                        spectra_ID = sii_to_spectra_id_and_ref[sii['id']][0]
                        spectra_ref = sii_to_spectra_id_and_ref[sii['id']][1]
                        pg_list.append((spectra_ref, pg_id, seq_accession, spectra_ID))
        df_pg = pd.DataFrame(pg_list, columns=["sample", "protein_group", "accession", "spectra_ID"])
        return df_pg

    def read_file_into_dataframe(self):
        # retrieve_refs = False: return IDs instead of all references
        reader = mzid.MzIdentML(source=self.result_table, iterative=True, use_index=True, retrieve_refs=False)
        spectra_dict = self.get_spectra_id_SpectrumIdentificationItem_dict(reader)
        # Reset the iterator, returning to the start of the file
        reader.reset()
        # sample | protein_group | accession | spectra_ID
        # qExactive01819.mgf  PAG_0             O75947        index=6865
        # qExactive01819.mgf  PAG_0             O75947        index=2422
        df_pg = self.create_pg_df(reader, spectra_dict)

        if df_pg.empty:
            raise AttributeError('Input error: ' + self.result_table + ' includes no protein ambiguity group information!\n'
                                                                       'Did you use mzident 1.1 format?')
        # "sample"  'accession'  'protein_group'    list of [spectra_ID]
        df_pg = df_pg.groupby(["sample", 'accession', 'protein_group'], as_index=False)\
            .agg({'spectra_ID': lambda x: sorted(list(x.unique()))})
        # "sample"  list of protein [accessions]  'protein_group'    list of list of [[spectra_IDs per accs],[]]
        df_pg = df_pg.groupby(["sample", 'protein_group'], as_index=False) \
            .agg({'accession': lambda x: list(x),
                  'spectra_ID': lambda x: list(x)})
        df_pg["quant"] =  df_pg['spectra_ID'].apply(lambda x: len(set([e for lst in x for e in lst])))
        # rename columns
        df_pg.columns = ["sample", 'protein_group', 'proteins', 'spectra_IDs', 'quant']
        # remove protein group column
        del df_pg['protein_group']
        return self.check_and_adapt_accessions(df_pg)
