# AUTOGENERATED! DO NOT EDIT! File to edit: nbdev_nbs/io/psm_reader/msfragger_reader.ipynb (unless otherwise specified).

__all__ = []

# Cell
import numpy as np
import pandas as pd
from alphabase.io.psm_reader.psm_reader import (
    PSMReaderBase, psm_reader_yaml,
    psm_reader_provider
)

from alphabase.constants.aa import AA_ASCII_MASS
from alphabase.constants.modification import MOD_INFO_DICT as mod_info

try:
    import pyteomics.pepxml as pepxml

    def _is_decoy(proteins):
        for prot in proteins:
            if not prot.startswith('rev_'):
                return False
        return True

    mass_mapped_mods = psm_reader_yaml['msfragger_pepxml']['mass_mapped_mods']
    mod_mass_tol = psm_reader_yaml['msfragger_pepxml']['mod_mass_tol']

    def _get_msf_mods(sequence, msf_aa_mods):
        mods = []
        mod_sites = []
        mod_deltas = []
        mod_delta_sites = []
        for mod in msf_aa_mods:
            mod_mass, site_str = mod.split('@')
            mod_mass = float(mod_mass)
            site = int(site_str)-1
            mod_mass = mod_mass - AA_ASCII_MASS[ord(sequence[site])]

            mod_considered = False
            for mod_name in mass_mapped_mods:
                if abs(mod_mass-mod_info[mod_name]['mass'])<mod_mass_tol:
                    if site == 0 and mod_name.endswith('N-term'):
                        mods.append(mod_name)
                        mod_sites.append('0')
                        mod_considered = True
                        break
                    _mod = mod_name.split('@')[0]+'@'+sequence[site]
                    if _mod in mod_info:
                        mods.append(_mod)
                        mod_sites.append(site_str)
                        mod_considered = True
                        break
            if not mod_considered:
                mod_deltas.append(str(mod_mass))
                mod_delta_sites.append(site_str)
        return ';'.join(mods), ';'.join(mod_sites), ';'.join(mod_deltas), ';'.join(mod_delta_sites)

    class MSFraggerPepXML(PSMReaderBase):
        def __init__(self, *,
            column_mapping: dict = None,
            modification_mapping: dict = None,
            keep_decoy=True,
            **kwargs
        ):
            """MSFragger is not fully supported as we can only access the pepxml file.
            """
            super().__init__(
                column_mapping=column_mapping,
                modification_mapping=modification_mapping,
                keep_decoy=keep_decoy,
                **kwargs
            )

        def _init_column_mapping(self):
            self.column_mapping = psm_reader_yaml[
                'msfragger_pepxml'
            ]['column_mapping']

        def _init_modification_mapping(self):
            self.modification_mapping = {}

        def _translate_modifications(self):
            pass

        def _load_file(self, filename):
            msf_df = pepxml.DataFrame(filename)
            msf_df.fillna('', inplace=True)
            msf_df.retention_time_sec /= 60
            msf_df['raw_name'] = msf_df[
                'spectrum'
            ].str.split('.').apply(lambda x: x[0])
            return msf_df

        def _translate_decoy(self, origin_df=None):
            self._psm_df['decoy'] = self._psm_df.proteins.apply(
                _is_decoy
            ).astype(np.int8)

            self._psm_df.proteins = self._psm_df.proteins.apply(
                lambda x: ';'.join(x)
            )
        def _translate_score(self, origin_df=None):
            if self.column_mapping['score'] == 'expect':
                # evalue score
                self._psm_df['score'] = -np.log(
                    self._psm_df['score']+1e-100
                )

        def _load_modifications(self, msf_df):
            (
                self._psm_df['mods'], self._psm_df['mod_sites'],
                self._psm_df['mod_deltas'], self._psm_df['mod_delta_sites'],
            ) = zip(*msf_df[['peptide','modifications']].apply(
                lambda x: _get_msf_mods(*x), axis=1)
            )

    psm_reader_provider.register_reader('msfragger_pepxml', MSFraggerPepXML)
except ModuleNotFoundError:
    pass
