# AUTOGENERATED! DO NOT EDIT! File to edit: nbdev_nbs/rescore/feature_extractor.ipynb (unless otherwise specified).

__all__ = ['match_one_raw', 'get_ms2_features', 'match_one_raw_mp', 'get_ms2_features_mp', 'perc_settings',
           'ScoreFeatureExtractor', 'ScoreFeatureExtractorMP']

# Cell
import pandas as pd
import numpy as np

from peptdeep.pretrained_models import ModelManager
from peptdeep.model.ms2 import calc_ms2_similarity
from peptdeep.mass_spec.match import PepSpecMatch
from alphabase.peptide.fragment import get_charged_frag_types
from alphabase.peptide.precursor import (
    refine_precursor_df
)
from alphabase.peptide.fragment import (
    concat_precursor_fragment_dataframes
)

from peptdeep.rescore.fdr import calc_fdr_for_df
from peptdeep.utils import process_bar, logging
from peptdeep.settings import global_settings
perc_settings = global_settings['percolator']


def match_one_raw(
    psm_df_one_raw,
    ms2_file,
    ms2_file_type,
    frag_types_to_match,
    ms2_ppm, ms2_tol,
):
    """ Internal function """
    match = PepSpecMatch(
        charged_frag_types=frag_types_to_match
    )

    # (
    #     psm_df, fragment_mz_df,
    #     matched_intensity_df, matched_mz_err_df
    # ) =
    return match.match_ms2_one_raw(
        refine_precursor_df(psm_df_one_raw),
        ms2_file=ms2_file,
        ms2_file_type=ms2_file_type,
        ppm=ms2_ppm, tol=ms2_tol,
    )

def get_ms2_features(
    psm_df, frag_types,
    predict_intensity_df,
    matched_intensity_df,
    matched_mass_err_df,
)->pd.DataFrame:
    """ Extract ms2 features from the given
    predict_intensity_df and matched_intensity_df. It will add columns into psm_df:
      cos: cosine similarity between predicted and matched fragments
      pcc: pearson correlation between predicted and matched fragments
      sa: spectral angle between predicted and matched fragments
      spc: Spearman's rank correlation between predicted and matched fragments.
      cos_bion: ...
      cos_yion: ...
      pcc_bion: ...
      pcc_yion: ...
      sa_bion: ...
      sa_yion: ...
      spc_bion: ...
      spc_yion: ...
      frag_ratio: # matched fragments / # total b&y fragments
      frag_ratio_bion: # matched b fragments / # total b fragments
      frag_ratio_yion: # matched y fragments / # total y fragments
    """
    used_frag_types = frag_types
    predict_intensity_df = predict_intensity_df[
        used_frag_types
    ]

    psm_df, ms2_metrics_df = calc_ms2_similarity(
        psm_df, predict_intensity_df,
        matched_intensity_df,
        charged_frag_types=used_frag_types,
        metrics=['COS','SA','SPC','PCC'],
        spc_top_k=perc_settings['top_k_frags_to_calc_spc']
    )
    psm_df.rename(
        columns={
            'COS':'cos','SA':'sa','SPC':'spc','PCC':'pcc',
        },
        inplace=True
    )

    b_frag_types = [
        _t for _t in used_frag_types
        if _t.startswith('b')
    ]
    y_frag_types = [
        _t for _t in used_frag_types
        if _t.startswith('y')
    ]

    frag_position_hits = matched_intensity_df[
        used_frag_types
    ].values.any(axis=1)
    frag_ratio_ion = []
    for start_idx, end_idx in psm_df[
        ['frag_start_idx','frag_end_idx']
    ].values:
        frag_ratio_ion.append(
            np.mean(frag_position_hits[start_idx:end_idx])
        )
    psm_df['frag_ratio'] = frag_ratio_ion

    if len(b_frag_types) > 0:
        psm_df, ms2_metrics_df = calc_ms2_similarity(
            psm_df, predict_intensity_df,
            matched_intensity_df,
            charged_frag_types=b_frag_types,
            metrics=['COS','SA','SPC','PCC'],
        )
        psm_df.rename(
            columns={
                'COS':'cos_bion','SA':'sa_bion','SPC':'spc_bion',
                'PCC':'pcc_bion'
            },
            inplace=True
        )
        frag_position_hits = matched_intensity_df[
            b_frag_types
        ].values.any(axis=1)

        frag_ratio_ion = []
        for start_idx, end_idx in psm_df[
            ['frag_start_idx','frag_end_idx']
        ].values:
            frag_ratio_ion.append(
                np.mean(frag_position_hits[start_idx:end_idx])
            )
        psm_df['frag_ratio_bion'] = frag_ratio_ion
    else:
        psm_df[['cos_bion','sa_bion','spc_bion','pcc_bion']] = 0
        psm_df[['frag_ratio_bion']] = 0

    if len(y_frag_types) > 0:
        psm_df, ms2_metrics_df = calc_ms2_similarity(
            psm_df, predict_intensity_df,
            matched_intensity_df,
            charged_frag_types=y_frag_types,
            metrics=['COS','SA','SPC'],
        )
        psm_df.rename(
            columns={
                'COS':'cos_yion','SA':'sa_yion','SPC':'spc_yion',
                'PCC':'pcc_yion',
            },
            inplace=True
        )
        frag_position_hits = matched_intensity_df[
            y_frag_types
        ].values.any(axis=1)
        frag_ratio_ion = []
        for start_idx, end_idx in psm_df[
            ['frag_start_idx','frag_end_idx']
        ].values:
            frag_ratio_ion.append(
                np.mean(frag_position_hits[start_idx:end_idx])
            )
        psm_df['frag_ratio_yion'] = frag_ratio_ion
    else:
        psm_df[['cos_yion','sa_yion','spc_yion','pcc_yion']] = 0
        psm_df[['frag_ratio_yion']] = 0

    return psm_df

# for imap/imap_unordered with multiprocessing.Pool()
def match_one_raw_mp(args):
    return match_one_raw(*args)

# for imap/imap_unordered with multiprocessing.Pool()
def get_ms2_features_mp(args):
    return get_ms2_features(*args)


# Cell

class ScoreFeatureExtractor:
    def __init__(self,
        model_mgr:ModelManager
    ):
        """ ScoreFeatureExtractor: Feature extractor for percolator
              with a single process.

        Args:
            model_mgr (ModelManager): The ModelManager in peptdeep.pretrained_models.
        """
        self.model_mgr = model_mgr
        self.model_mgr.verbose = False

        self.raw_num_to_tune = perc_settings['raw_num_to_tune']

        (
            self.model_mgr.psm_num_to_tune_ms2
        ) = perc_settings['psm_num_to_tune_ms2']

        (
            self.model_mgr.psm_num_per_mod_to_tune_ms2
        ) = perc_settings['psm_num_per_mod_to_tune_ms2']

        (
            self.model_mgr.psm_num_to_tune_rt_ccs
        ) = perc_settings['psm_num_to_tune_rt_ccs']

        (
            self.model_mgr.mod_psm_num_to_tune_rt_ccs
        ) = perc_settings['mod_psm_num_to_tune_rt_ccs']

        (
            self.model_mgr.top_n_mods_to_tune
        ) = perc_settings['top_n_mods_to_tune']

        self.require_model_tuning = perc_settings[
            'require_model_tuning'
        ]
        self.require_raw_specific_tuning = perc_settings[
            'require_raw_specific_tuning'
        ]
        self.raw_specific_ms2_tuning = perc_settings[
            'raw_specific_ms2_tuning'
        ]

        self.score_feature_list = [
            'sa','spc','pcc',
            'sa_bion','spc_bion','pcc_bion',
            'sa_yion','spc_yion','pcc_yion',
            'frag_ratio','frag_ratio_bion',
            'frag_ratio_yion','rt_delta_abs',
            'mobility_delta_abs',
        ]

    def _select_raw_to_tune(self,
        psm_df:pd.DataFrame,
    )->tuple:
        """ Randomly select `self.raw_num_to_tune` raw files
        to tune the models. If # raw files is less than `self.raw_num_to_tune`,
        all raw files will be used to tune the model.

        Args:
            psm_df (pd.DataFrame): dataframe contains PSMs of all raw files.

        Returns:
            df_groupby_raw: psm_df.groupby('raw_name')
            list: selected raw_name list
        """
        if 'fdr' not in psm_df.columns:
            psm_df = calc_fdr_for_df(psm_df, 'score')
        df_fdr = psm_df[(psm_df.fdr<0.01)&(psm_df.decoy==0)]

        df_groupby_raw = df_fdr.groupby('raw_name')

        if df_groupby_raw.ngroups < self.raw_num_to_tune:
            tune_raw_num = df_groupby_raw.ngroups
        else:
            tune_raw_num = self.raw_num_to_tune

        raw_list = list(
            df_groupby_raw['score'].count().rank(
                ascending=True
            ).nlargest(tune_raw_num).index
        )

        return df_groupby_raw, raw_list

    def fine_tune_models(self,
        psm_df:pd.DataFrame,
        ms2_file_dict:dict,
        ms2_file_type:str,
        frag_types_to_match:str,
        ms2_ppm:bool, ms2_tol:float,
    ):
        """ Extract information from the ms2 files,
        and then fine-tune the models

        Args:
            psm_df (pd.DataFrame): psm_df
            ms2_file_dict (dict): {raw_name: ms2_file_path}
            ms2_file_type (str): ms2_file_type, could be 'alphapept', 'mgf', 'thermo_raw'
            frag_types_to_match (str): ['b_z1','b_z2','y_z1'...]
            ms2_ppm (bool): is ppm tolerance for ms2 matching
            ms2_tol (float): tolerance value for ms2 matching
        """
        logging.info('Preparing for fine-tuning ...')

        (
            df_groupby_raw, raw_list
        ) = self._select_raw_to_tune(psm_df)

        psm_df_list = []
        matched_intensity_df_list = []
        for raw_name, df in process_bar(
            df_groupby_raw, df_groupby_raw.ngroups
        ):
            if (
                raw_name not in raw_list
                or raw_name not in ms2_file_dict
            ):
                continue
            (
                df, _, inten_df, _
            ) = match_one_raw(
                df, ms2_file_dict[raw_name],
                ms2_file_type,
                frag_types_to_match,
                ms2_ppm, ms2_tol,
            )
            psm_df_list.append(df)
            matched_intensity_df_list.append(inten_df)

        logging.info('Fine-tuning ...')
        if len(psm_df_list) == 0: return
        self._tune(
            *concat_precursor_fragment_dataframes(
                psm_df_list,
                matched_intensity_df_list
            )
        )

    def _tune(self,
        psm_df,
        matched_intensity_df
    ):
        if 'ccs' in psm_df.columns:
            self.model_mgr.fine_tune_ccs_model(psm_df)

        self.model_mgr.fine_tune_rt_model(psm_df)

        self.model_mgr.fine_tune_ms2_model(
            psm_df, matched_intensity_df
        )

    def extract_rt_features(self, psm_df):
        if self.require_raw_specific_tuning:
            (
                psm_num_to_tune_rt_ccs,
                mod_psm_num_to_tune_rt_ccs,
                epoch_to_tune_rt_ccs
            ) = (
                self.model_mgr.psm_num_to_tune_rt_ccs,
                self.model_mgr.mod_psm_num_to_tune_rt_ccs,
                self.model_mgr.epoch_to_tune_rt_ccs
            )

            (
                self.model_mgr.psm_num_to_tune_rt_ccs
            ) = perc_settings['psm_num_per_raw_to_tune']

            self.model_mgr.mod_psm_num_to_tune_rt_ccs = 0

            (
                self.model_mgr.epoch_to_tune_rt_ccs
            ) = perc_settings['epoch_per_raw_to_tune']
            self.model_mgr.fine_tune_rt_model(
                psm_df[(psm_df.fdr<0.01)&(psm_df.decoy==0)]
            )

            (
                self.model_mgr.psm_num_to_tune_rt_ccs,
                self.model_mgr.mod_psm_num_to_tune_rt_ccs,
                self.model_mgr.epoch_to_tune_rt_ccs
            ) = (
                psm_num_to_tune_rt_ccs,
                mod_psm_num_to_tune_rt_ccs,
                epoch_to_tune_rt_ccs
            )

        if 'rt_norm' in psm_df.columns:
            psm_df = self.model_mgr.predict_rt(
                psm_df
            )
            psm_df[
                'rt_delta'
            ] = (
                psm_df.rt_pred-psm_df.rt_norm
            )

            mean_delta = psm_df.loc[
                (psm_df.fdr<0.01)&(psm_df.decoy==0),
                'rt_delta'
            ].mean()

            if np.isnan(mean_delta):
                mean_delta = 0

            psm_df['rt_delta_abs'] = (
                psm_df.rt_delta-mean_delta
            ).abs()
        else:
            psm_df['rt_delta'] = 0
            psm_df['rt_delta_abs'] = 0

    def extract_mobility_features(self, psm_df):
        if (
            'mobility' in psm_df.columns
        ):
            psm_df = self.model_mgr.predict_mobility(
                psm_df
            )

            psm_df[
                'mobility_delta'
            ] = (
                psm_df.mobility_pred-psm_df.mobility
            )

            mean_delta = psm_df.loc[
                (psm_df.fdr<0.01)&(psm_df.decoy==0),
                'mobility_delta'
            ].mean()

            if np.isnan(mean_delta):
                mean_delta = 0

            psm_df['mobility_delta_abs'] = (
                psm_df.mobility_delta-mean_delta
            ).abs()
        else:
            psm_df['mobility_delta'] = 0
            psm_df['mobility_delta_abs'] = 0


    def match_ms2(self,
        psm_df: pd.DataFrame,
        ms2_file_dict, #raw_name: ms2_file_path or ms_reader object
        ms2_file_type:str,
        frag_types_to_match:list = get_charged_frag_types(['b','y'], 2),
        ms2_ppm=True, ms2_tol=20,
    ):
        self.match = PepSpecMatch(
            charged_frag_types=frag_types_to_match
        )

        self.match.match_ms2_centroid(
            refine_precursor_df(psm_df),
            ms2_file_dict=ms2_file_dict,
            ms2_file_type=ms2_file_type,
            ppm=ms2_ppm, tol=ms2_tol,
        )

    def _get_model_frag_types(self, frag_types):
        used_frag_types = []
        for frag_type in frag_types:
            if frag_type in (
                self.model_mgr.ms2_model.charged_frag_types
            ):
                used_frag_types.append(frag_type)
        return used_frag_types

    def extract_features(self,
        psm_df: pd.DataFrame,
        ms2_file_dict,
        ms2_file_type,
        frag_types:list = get_charged_frag_types(['b','y'], 2),
        ms2_ppm=True, ms2_tol=20,
    )->pd.DataFrame:
        """ Extract features and add columns (self.score_feature_list) into psm_df

        Args:
            psm_df (pd.DataFrame): psm dataframe to extract features
            ms2_file_dict ([type]): MS2 file path dict: {raw_name: ms2_path}
            ms2_file_type (str, optional): MS2 file type, coult be
              'alphapept', 'mgf', or 'raw'.
            frag_types (list, optional): fragment types.
              Defaults to `alphabase.fragment.get_charged_frag_types(['b','y'], 2)`.
            ms2_ppm (bool, optional): Matching MS2 mass tolerance unit.
              Defaults to True.
            ms2_tol (int, optional): Matching mass tolerance.
              Defaults to 20.

        Returns:
            pd.DataFrame: psm_df with feature columns added
        """

        frag_types = self._get_model_frag_types(frag_types)

        if self.require_model_tuning:
            logging.info('Fine-tuning models ...')
            self.fine_tune_models(
                psm_df,
                ms2_file_dict, ms2_file_type,
                frag_types, ms2_ppm, ms2_tol
            )

        logging.info('Extracting peptdeep features ...')
        result_psm_list = []
        groupby = psm_df.groupby('raw_name')
        for raw_name, df in process_bar(groupby, groupby.ngroups):
            if raw_name not in ms2_file_dict:
                continue
            (
                df, frag_mz_df, frag_inten_df, frag_merr_df
            ) = match_one_raw(
                df,
                ms2_file_dict[raw_name],
                ms2_file_type,
                frag_types,
                ms2_ppm, ms2_tol,
            )

            self.extract_rt_features(df)
            self.extract_mobility_features(df)

            predict_inten_df = self.model_mgr.predict_ms2(df)

            result_psm_list.append(
                get_ms2_features(
                    df, frag_types,
                    predict_inten_df,
                    frag_inten_df,
                    frag_merr_df,
                )
            )

        self.psm_df = pd.concat(
            result_psm_list
        ).reset_index(drop=True)
        logging.info('Finish extracting features')
        return self.psm_df


# Cell
import torch
import torch.multiprocessing as mp


class ScoreFeatureExtractorMP(ScoreFeatureExtractor):
    def __init__(self,
        model_mgr:ModelManager
    ):
        """ ScoreFeatureExtractorMP: Feature extractor for percolator
              with multiprocessing.

        Args:
            model_mgr (ModelManager): The ModelManager in peptdeep.pretrained_models.
        """
        super().__init__(model_mgr=model_mgr)

        # share_memory to save memory
        self.model_mgr.ms2_model.model.share_memory()
        self.model_mgr.rt_model.model.share_memory()
        self.model_mgr.ccs_model.model.share_memory()


    def fine_tune_models(self,
        psm_df,
        ms2_file_dict,
        ms2_file_type,
        frag_types_to_match,
        ms2_ppm, ms2_tol,
    ):
        """ Extract information from the ms2 files,
        and then fine-tune the models

        Args:
            psm_df (pd.DataFrame): psm_df
            ms2_file_dict (dict): {raw_name: ms2_file_path}
            ms2_file_type (str): ms2_file_type, could be 'alphapept', 'mgf', 'thermo_raw'
            frag_types_to_match (str): ['b_z1','b_z2','y_z1'...]
            ms2_ppm (bool): is ppm tolerance for ms2 matching
            ms2_tol (float): tolerance value for ms2 matching
        """
        (
            df_groupby_raw, raw_list
        ) = self._select_raw_to_tune(psm_df)

        def one_raw_param_generator(df_groupby_raw):
            for raw_name, df in df_groupby_raw:
                if (
                    raw_name not in raw_list
                    or raw_name not in ms2_file_dict
                ):
                    continue

                yield (
                    df,
                    ms2_file_dict[raw_name],
                    ms2_file_type,
                    frag_types_to_match,
                    ms2_ppm, ms2_tol,
                )

        logging.info('Preparing for fine-tuning ...')
        psm_df_list = []
        matched_intensity_df_list = []
        with mp.Pool(global_settings['thread_num']) as p:
            for df, _, inten_df, _ in process_bar(
                p.imap_unordered(
                    match_one_raw_mp,
                    one_raw_param_generator(df_groupby_raw)
                ), df_groupby_raw.ngroups
            ):
                psm_df_list.append(df)
                matched_intensity_df_list.append(inten_df)

        logging.info('Fine-tuning ...')
        if len(psm_df_list) == 0: return
        self._tune(
            *concat_precursor_fragment_dataframes(
                psm_df_list,
                matched_intensity_df_list
            )
        )

    def extract_features_one_raw(self,
        df_one_raw: pd.DataFrame,
        ms2_file,
        ms2_file_type,
        frag_types,
        ms2_ppm, ms2_tol,
    ):
        (
            df, frag_mz_df, frag_inten_df, frag_merr_df
        ) = match_one_raw(df_one_raw,
            ms2_file, ms2_file_type, frag_types,
            ms2_ppm, ms2_tol
        )

        self.extract_rt_features(df)
        self.extract_mobility_features(df)

        predict_inten_df = self.model_mgr.predict_ms2(df)

        return get_ms2_features(df,
            frag_types,
            predict_inten_df,
            frag_inten_df,
            frag_merr_df,
        )

    def extract_features(self,
        psm_df: pd.DataFrame,
        ms2_file_dict,
        ms2_file_type,
        frag_types:list = get_charged_frag_types(['b','y'], 2),
        ms2_ppm=True, ms2_tol=20,
    )->pd.DataFrame:

        """ MPExtract features and add columns (self.score_feature_list) into psm_df.
        Main entry of ScoreFeatureExtractor.

        Args:
            psm_df (pd.DataFrame): psm dataframe to extract features
            ms2_file_dict ([type]): MS2 file path dict: {raw_name: ms2_path}
            ms2_file_type (str, optional): MS2 file type, coult be
              'alphapept', 'mgf', or 'thermo'.
            frag_types (list, optional): fragment types.
              Defaults to `alphabase.fragment.get_charged_frag_types(['b','y'], 2)`.
            ms2_ppm (bool, optional): Matching MS2 mass tolerance unit.
              Defaults to True.
            ms2_tol (int, optional): Matching mass tolerance.
              Defaults to 20.

        Returns:
            pd.DataFrame: psm_df with feature columns added
        """

        used_frag_types = self._get_model_frag_types(frag_types)

        if self.require_model_tuning:
            logging.info('Fine-tuning models ...')
            self.fine_tune_models(
                psm_df,
                ms2_file_dict, ms2_file_type,
                used_frag_types, ms2_ppm, ms2_tol
            )

        def one_raw_param_generator(df_groupby_raw):
            for raw_name, df in df_groupby_raw:
                if raw_name not in ms2_file_dict:
                    continue
                yield (
                    df,
                    ms2_file_dict[raw_name],
                    ms2_file_type,
                    used_frag_types,
                    ms2_ppm, ms2_tol,
                )

        logging.info(
            'Extracting peptdeep features with multiprocessing ...'
        )
        df_groupby_raw = psm_df.groupby('raw_name')
        result_psm_list = []
        if (
            not torch.cuda.is_available() and
            not self.require_raw_specific_tuning
        ):
            # use multiprocessing for prediction
            # only when no GPUs are available
            with mp.Pool(global_settings['thread_num']) as p:
                for _df in process_bar(p.imap_unordered(
                    self.extract_features_one_raw,
                    one_raw_param_generator(df_groupby_raw)
                ), df_groupby_raw.ngroups):
                    result_psm_list.append(_df)

        else:
            # multiprocessing is only used for ms2 matching
            def prediction_gen(df_groupby_raw):
                with mp.Pool(global_settings['thread_num']) as _p:
                    for (
                        df, frag_mz_df, frag_inten_df, frag_merr_df
                    ) in _p.imap_unordered(
                        match_one_raw_mp,
                        one_raw_param_generator(df_groupby_raw)
                    ):
                        # outsite multiprocessing region
                        self.extract_rt_features(df)
                        self.extract_mobility_features(df)

                        if (
                            self.require_raw_specific_tuning
                            and self.raw_specific_ms2_tuning
                        ):
                            (
                                psm_num_to_tune_ms2,
                                psm_num_per_mod_to_tune_ms2,
                                epoch_to_tune_ms2,
                                use_grid_nce_search
                            ) = (
                                self.model_mgr.psm_num_to_tune_ms2,
                                self.model_mgr.psm_num_per_mod_to_tune_ms2,
                                self.model_mgr.epoch_to_tune_ms2,
                                self.model_mgr.use_grid_nce_search
                            )

                            (
                                self.model_mgr.psm_num_to_tune_ms2
                            ) = perc_settings['psm_num_per_raw_to_tune']

                            self.model_mgr.psm_num_per_mod_to_tune_ms2 = 0

                            self.model_mgr.epoch_to_tune_ms2 = 3

                            self.model_mgr.use_grid_nce_search = False

                            if 'nce' not in df.columns:
                                self.model_mgr.set_default_nce(df)

                            self.model_mgr.fine_tune_ms2_model(
                                df[(df.fdr<0.01)&(df.decoy==0)],
                                frag_inten_df
                            )

                            (
                                self.model_mgr.psm_num_to_tune_ms2,
                                self.model_mgr.psm_num_per_mod_to_tune_ms2,
                                self.model_mgr.epoch_to_tune_ms2,
                                self.model_mgr.use_grid_nce_search
                            ) = (
                                psm_num_to_tune_ms2,
                                psm_num_per_mod_to_tune_ms2,
                                epoch_to_tune_ms2,
                                use_grid_nce_search
                            )

                        predict_inten_df = self.model_mgr.predict_ms2(df)

                        yield (
                            df, used_frag_types,
                            predict_inten_df,
                            frag_inten_df, frag_merr_df,
                        )

            with mp.Pool(global_settings['thread_num']) as p:
                for df in process_bar(p.imap_unordered(
                    get_ms2_features_mp,
                    prediction_gen(df_groupby_raw)
                ), df_groupby_raw.ngroups):
                    result_psm_list.append(df)

        self.psm_df = pd.concat(
            result_psm_list
        ).reset_index(drop=True)
        logging.info('Finished feature extraction with multiprocessing')
        return self.psm_df