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

import os
import sys
import xml.etree.ElementTree as ET
import pandas as pd
from utils.search_result_parsing.helper_classes import ProteomicSearchResultReaderBaseClass

sys.path.append(os.getcwd())


class MzidentTableReader(ProteomicSearchResultReaderBaseClass):
    """
    params: ProteomicSearchResultReaderBaseClass = BaseClass
        params base class: result_table (from BaseClass ProteomicSearchResultParser) = mzid file to read in
        params base class: f_style = yaml (mzTab_1_1_0.yaml) (from BaseClass ProteomicSearchResultParser)
        params base class: required_sample_descriptors (from BaseClass ProteomicSearchResultParser) = '1'
        params base class: self.style_dict = load_yaml(self.f_style)
    """

    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 db_sequence_df_def(self, xml_root, name_space):
            """
    :param xml_root:
    :param name_space:
    :return:
    Fetches protein (dBSequence_ref) information from mzid file
    accession and id are included in tag <DBSequence>

    Args:
        xml_root: an xml root object
        name_space: xml files may indicate name space
    Return:
        a pandas data frame includes DBSequence information.
        DF fields: dBSequence_ref and accession
    """
            DBSequence_list = []
            for DBSequence in xml_root.iter(name_space + 'DBSequence'):
                DBSequence_item = {}
                DBSequence_item['dBSequence_ref'] = DBSequence.get('id')
                DBSequence_item['accession'] = DBSequence.get('accession')
                DBSequence_list = DBSequence_list + [DBSequence_item]
            return pd.DataFrame(DBSequence_list)

    def db_sequence_df_def(self, xml_root, name_space):
        """
        :param xml_root:
        :param name_space:
        :return:
        Fetches protein (dBSequence_ref) information from mzid file
        accession and id are included in tag <DBSequence>

        Args:
            xml_root: an xml root object
            name_space: xml files may indicate name space
        Return:
            a pandas data frame includes DBSequence information.
            DF fields: dBSequence_ref and accession
        """
        DBSequence_list = []
        for DBSequence in xml_root.iter(name_space + 'DBSequence'):
            DBSequence_item = {}
            DBSequence_item['dBSequence_ref'] = DBSequence.get('id')
            DBSequence_item['accession'] = DBSequence.get('accession')
            DBSequence_list = DBSequence_list + [DBSequence_item]
        return pd.DataFrame(DBSequence_list)

    # TODO insert: leading protein, representative protein, test: confidence, pass threshold
    def proteinAmbiguityGroup_df_def(self, xml_root, name_space):
        """Fetches protein group with spectrum information from <ProteinAmbiguityGroup>

    Args:
        xml_root: an xml root object
        name_space: xml files may indicate name space
    Return:
        a pandas data frame includes protein group and spectrum information (one row per spectrum)
        DF fields - protein_group, dBSequence_ref and spectrum_identification
    """

        def readPAG(proteinAmbiguityGroup_list, pdh):
            dBSequence = pdh.attrib['dBSequence_ref']
            for spectrum_identification_item in pdh.iter(name_space+'SpectrumIdentificationItemRef'):
                proteinAmbiguityGroup_dict = {}
                proteinAmbiguityGroup_dict['protein_group'] = proteinAmbiguityGroup.attrib['id']#.split('_')[1]
                proteinAmbiguityGroup_dict['dBSequence_ref'] = dBSequence
                proteinAmbiguityGroup_dict['spectrum_identification'] = spectrum_identification_item.attrib['spectrumIdentificationItem_ref']
                proteinAmbiguityGroup_list.append(proteinAmbiguityGroup_dict)

        def fill_df(proteinAmbiguityGroup_list, proteinAmbiguityGroup):
            for protein_detection_hypothesis in proteinAmbiguityGroup.iter(name_space+'ProteinDetectionHypothesis'):
                try:
                    if protein_detection_hypothesis.attrib['passThreshold'] == 'true':
                        readPAG(proteinAmbiguityGroup_list, protein_detection_hypothesis)
                # value passThreshold MAY be contained
                except KeyError:
                    readPAG(proteinAmbiguityGroup_list, protein_detection_hypothesis)

        # iter protein ambiguity groups, list of dict one per PAG
        proteinAmbiguityGroup_list = []
        threshold = False
        for proteinAmbiguityGroup in xml_root.iter(name_space + 'ProteinAmbiguityGroup'):
            for cvParam in proteinAmbiguityGroup.iter(name_space+'cvParam'):
                # cvParam MS:1002415 =  “protein group passes threshold” MUST be contained (but is not always)
                if "MS:1002415" in cvParam.attrib['accession']:
                    threshold = True
            if threshold:
                for cvParam in proteinAmbiguityGroup.iter(name_space+'cvParam'):
                    if cvParam.attrib['accession'] == "MS:1002415" and cvParam.attrib['value'] == 'true':
                        fill_df(proteinAmbiguityGroup_list, proteinAmbiguityGroup)
            else:
                fill_df(proteinAmbiguityGroup_list, proteinAmbiguityGroup)
        return pd.DataFrame(proteinAmbiguityGroup_list)


    def sample_spectrum_df_def(self, xml_root, name_space):
        '''Fetches spectrum and its coresponding sample info
    In some cases, <AnalysisSampleCollection> tag exist, while it is
    not indicated in most of the time. But it is still possible to get
    the information form <SpectrumIdentificationResult>, spectraData_ref
    is the sample id.

    Args:
        xml_root: an xml root object
        name_space: xml files may indicate name space
    Return:
        a pandas data frame includes sample and stpectrum
        DF fields: sample, spectrum_identification
        only if passThreshold="true"
    '''
        sample_spectrum_list = []
        for spectrumIdentification in xml_root.iter(name_space + 'SpectrumIdentificationResult'):
            sample_spectrum = {}
            for child in spectrumIdentification:
                try:
                    if child.attrib['passThreshold'] == 'true':
                        sample_spectrum['sample'] = spectrumIdentification.get('spectraData_ref')
                        spectrumIdentification_item = spectrumIdentification.find(name_space + 'SpectrumIdentificationItem')
                        sample_spectrum['spectrum_identification'] = spectrumIdentification_item.get('id')
                        sample_spectrum_list = sample_spectrum_list + [sample_spectrum]
                    # value passThreshold MAY be contained
                except KeyError:
                    sample_spectrum['sample'] = spectrumIdentification.get('spectraData_ref')
                    spectrumIdentification_item = spectrumIdentification.find(name_space + 'SpectrumIdentificationItem')
                    sample_spectrum['spectrum_identification'] = spectrumIdentification_item.get('id')
                    sample_spectrum_list = sample_spectrum_list + [sample_spectrum]
        return pd.DataFrame(sample_spectrum_list)

    def read_mizd_def(self, mzid_file):
        '''Reads mzid file
        Args:
            mzid_file: a mzid file
        Return:
            an xml root object read from mzid file.
            param root= ET.parse(mzid_file).tree.getroot() = XML data imported
            param namespace: xml namespace = xmlns entry in root tag eg.: in <MzIdentML ... xmlns={http://psidev.info/psi/pi/mzIdentML/1.2}>
        Raises:
            raises errors when file not found or no group info
        '''
        try:
            # read mzid file
            tree = ET.parse(mzid_file)
            root = tree.getroot()
        except FileNotFoundError:
            raise RuntimeError(mzid_file + ' can not be found!')

        # namespace
        if root.tag[0] == '{':
            name_space = '{' + root.tag.split('}')[0].strip('{}') + '}'
        else:
            name_space = ''

        # <ProteinDetectionList id="Protein_groups"> id=unique name in scope
        # test if <ProteinDetectionList> not exists, and raise error
        if not any(True for _ in root.iter(name_space + 'ProteinDetectionList')):
            # return 'There is no protein group information.'
            raise AttributeError(mzid_file + ' includes no protein group information!')

        # check if protein ambiguous groups contained = mzident format 1.2
        if not any(True for _ in root.iter(name_space + 'ProteinAmbiguityGroup')):
            # return 'There is no protein group information.'
            raise AttributeError('Input error: ' + mzid_file + ' includes no protein ambiguity group information!\n'
                                                               'Did you use mzident 1.1 format?')
        return root, name_space

    def read_file_into_dataframe(self):
        '''Fetch protein accessions (one group) and its spectrum count
        Args:
            mzid_file: a string to indicate mzid file
        Return:
            a data frame includes protein accessions (one group) and its spectrum count and spectra IDs per accession
        '''
        # root = read xml file, exit if no ProteinList contained or no protein ambig groups
        root, name_space = self.read_mizd_def(self.result_table)

        # dataframe with sequence references and accessions
        #  dBSequence_ref        accession
        # 0           P27348           P27348
        # 1  P27348_REVERSED  P27348_REVERSED
        df_seq = self.db_sequence_df_def(root, name_space)

        # # dataframe protein ambiguity protein_group, seq_ref and identified spectra (per spectrum one entry)
        #        protein_group   dBSequence_ref spectrum_identification
        # 0       #  PAG_0           O75947              SII_5365_1
        # 1       #  PAG_0           O75947              SII_6246_1
        df_protein = self.proteinAmbiguityGroup_df_def(root, name_space)

        # dataframe sample and spectrum_identification
        #         sample spectrum_identification
        # 0  qExactive01819.mgf                 SII_1_1
        # 1  qExactive01819.mgf                 SII_2_1
        df_spectrum = self.sample_spectrum_df_def(root, name_space)

        # sample spectrum_identification protein_group   dBSequence_ref
        # 0  qExactive01819.mgf                 SII_1_1      PAG_1583  Q13428_REVERSED
        # 1  qExactive01819.mgf                 SII_2_1      PAG_1661           O95456
        sample_spectrum_group_dBSequence_df = pd.merge(df_spectrum, df_protein, on='spectrum_identification', how='inner')

        #                sample spectrum_identification protein_group dBSequence_ref    accession
        # 0      qExactive2.mgf              SII_1769_1      PAG_2155         P21333   P21333
        # 1      qExactive2.mgf              SII_1769_1      PAG_2155         P21333   P21333
        # 2      qExactive2.mgf              SII_1769_1      PAG_2155         P21333   P21333
        sample_spectrum_group_peptide_accession_df = pd.merge(sample_spectrum_group_dBSequence_df, df_seq, on='dBSequence_ref', how='inner')

        sample_acc_spectra_dict = {}
        # [sample][protein] = spectra list
        for i, row in sample_spectrum_group_peptide_accession_df.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['spectrum_identification'])
                else:
                    sample_acc_spectra_dict[row['sample']][str(row['accession'])] = {row['spectrum_identification']}
            else:
                sample_acc_spectra_dict[row['sample']] = {str(row['accession']): {row['spectrum_identification']}}

        group_accessions_df = sample_spectrum_group_peptide_accession_df[['protein_group','accession']].groupby('protein_group', as_index=False) \
            .agg({'accession': lambda x: list(x.unique())})

        spectra_per_protein = {}
        for sample in sample_acc_spectra_dict.keys():
            spectra_per_protein[sample] = {}
            for proteins in group_accessions_df['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:
                            spectra_per_protein[sample][protein] = sorted(sample_acc_spectra_dict[sample][protein])
                        except KeyError:
                            spectra_per_protein[sample][protein] = ['']

        # build df from lists
        quant = []
        spectra_IDs = []
        proteins = []
        samples = []
        for sample in sample_acc_spectra_dict.keys():
            for accessions in group_accessions_df['accession'].tolist():
                proteins_spectra = []
                for protein in accessions:
                    try:
                        proteins_spectra.append(list(spectra_per_protein[sample][protein]))
                    except (KeyError, AttributeError):
                        continue
                if proteins_spectra:
                    spectra_IDs.append(proteins_spectra)
                    proteins.append(accessions)
                    samples.append(sample)
                    quant.append(len(set([val for sublist in proteins_spectra for val in sublist])))
        df_standard = pd.DataFrame({'sample': samples, 'proteins': proteins, 'quant':quant, 'spectra_IDs': spectra_IDs})

        return self.check_and_adapt_accessions(df_standard)
