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

__all__ = ['get_reduced_mass', 'ccs_to_mobility_bruker', 'mobility_to_ccs_bruker', 'ccs_to_mobility_for_df',
           'mobility_to_ccs_for_df', 'CCS_IM_COEF', 'IM_GAS_MASS']

# Cell
import numpy as np
import pandas as pd

from alphabase.peptide.fragment import update_precursor_mz
from alphabase.constants.element import common_const_dict

CCS_IM_COEF = common_const_dict['MOBILITY']['CCS_IM_COEF']
IM_GAS_MASS = common_const_dict['MOBILITY']['IM_GAS_MASS']

def get_reduced_mass(
    precursor_mzs: np.array,
    charges: np.array
)->np.array:
    """ Reduced mass for CCS and mobility calculation """
    reduced_masses = precursor_mzs*charges
    return reduced_masses*IM_GAS_MASS/(reduced_masses+IM_GAS_MASS)

def ccs_to_mobility_bruker(
    ccs_values: np.array,
    charges: np.array,
    precursor_mzs: np.array
)->np.array:
    """ Convert CCS to mobility for Bruker (timsTOF) """
    reduced_masses = get_reduced_mass(precursor_mzs, charges)
    return ccs_values*np.sqrt(reduced_masses)/charges/CCS_IM_COEF

def mobility_to_ccs_bruker(
    im_values: np.array,
    charges: np.array,
    precursor_mzs: np.array
)->np.array:
    """ Convert mobility to CCS for Bruker (timsTOF) """
    reduced_masses = get_reduced_mass(precursor_mzs, charges)
    return im_values*charges*CCS_IM_COEF/np.sqrt(reduced_masses)

def ccs_to_mobility_for_df(
    precursor_df:pd.DataFrame,
    ccs_column:str,
    *,
    vendor="bruker"
)->np.array:
    """

    Args:
        precursor_df (pd.DataFrame): precursor_df
        ccs_column (str): CCS column name in precursor_df
        vendor (str, optional): Different vender may have different IM calculation.
            Defaults to "bruker".
            Note that other vendors are not implemented currently.

    Returns:
        np.array: mobility values
    """
    if 'precursor_mz' not in precursor_df.columns:
        precursor_df = update_precursor_mz(precursor_df)
    return ccs_to_mobility_bruker(
        precursor_df[ccs_column].values,
        precursor_df.charge.values,
        precursor_df.precursor_mz.values
    )

def mobility_to_ccs_for_df(
    precursor_df:pd.DataFrame,
    mobility_column:str,
    *,
    vendor="bruker"
)->np.array:
    """

    Args:
        precursor_df (pd.DataFrame): precursor_df
        mobility_column (str): mobility column name in precursor_df
        vendor (str, optional): Different vender may have different IM calculation.
            Defaults to "bruker".
            Note that other vendors are not implemented currently.

    Returns:
        np.array: CCS values
    """

    if 'precursor_mz' not in precursor_df.columns:
        precursor_df = update_precursor_mz(precursor_df)
    return mobility_to_ccs_bruker(
        precursor_df[mobility_column].values,
        precursor_df.charge.values,
        precursor_df.precursor_mz.values
    )