# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
# default_exp matching
# -

# # Matching
#
# > Functions related to matching

#hide
from nbdev.showdoc import *

# ## Dataset Alignment
#
# For matching MS2-identifications to MS1-features, we first need to align the datasets on top of each other to be able to transfer identifications correctly. Datasets are aligned by comparing shared precursors and calculating the median offset.
# When comparing all files to each other, we get an overdetermined linear equation system. By solving this, we find offset parameters that minimize the shift of all files to each other.
# Offset is either applied relative (mz, mobility) or absolute (rt).
#
# ### Relative offset
#
# For some parameters, we would like to have a relative correction of values. Consider the case of different `mz`-values, e.g. 300 and 600. If we assume that the offset is larger for larger m/z values, we would not want an absolute correction of e.g. +0.5 Da (300.5 and 600.5) but rather a relative correction of e.g. +0.1% (300.3 and 600.6).
#
# ### Absolute correction
#
# In contrast to the relative correction, sometimes absolute correction is more applicable. Consider the case of retention time. Here one would rather not expect a relative offset but rather an absolute offset. As an example, consider a lag time of 0.5 Minutes. This would be constant for all retention times and not differ e.g., for later retention times.

# +
#export

import pandas as pd
import numpy as np

def calculate_distance(table_1: pd.DataFrame, table_2: pd.DataFrame, offset_dict: dict, calib: bool = False) -> (list, int):
    """Calculate the distance between two precursors for different columns
    Distance can either be relative or absolute.
    
    An example for a minimal offset_dict is: offset_dict = {'mass':'absolute'}

    Args:
        table_1 (pd.DataFrame): Dataframe with precusor data.
        table_2 (pd.DataFrame): Dataframe with precusor data.
        offset_dict (dict): Dictionary with column names and how the distance should be calculated.
        calib (bool): Flag to indicate that distances should be calculated on calibrated columns. Defaults to False.
        
    Raises:
        KeyError: If either table_1 or table_2 is not indexed by precursor
        
    """

    if table_1.index.name != 'precursor':
        raise KeyError('table_1 is not indexed by precursor')

    if table_2.index.name != 'precursor':
        raise KeyError('table_2 is not indexed by precursor')

    shared_precursors = list(set(table_1.index).intersection(set(table_2.index)))

    table_1_ = table_1.loc[shared_precursors]
    table_2_ = table_2.loc[shared_precursors]

    table_1_ = table_1_.groupby('precursor').mean()
    table_2_ = table_2_.groupby('precursor').mean()

    deltas = []

    for col in offset_dict:
        if calib:
            col_ = col+'_calib'
        else:
            col_ = col

        if offset_dict[col] == 'absolute':
            deltas.append(np.nanmedian(table_1_[col_] - table_2_[col_]))
        elif offset_dict[col] == 'relative':
            deltas.append(np.nanmedian((table_1_[col_] - table_2_[col_]) / (table_1_[col_] + table_2_[col_]) * 2))
        else:
            raise NotImplementedError(f"Calculating delta for {offset_dict[col_]} not implemented.")

    return deltas, len(shared_precursors)


# +
#hide
def test_calcualte_distance():
    a = pd.DataFrame({'precursor':['A','B','C','D','E'], 'mass':[10, 20, 30, 40, 50]}).set_index('precursor')
    b = pd.DataFrame({'precursor':['A','B','C','D','F'], 'mass':[11, 21, 31, 41, 51]}).set_index('precursor')
    offset_dict = {'mass':'absolute'}

    delta, n_shared = calculate_distance(a, a, offset_dict)
    assert (delta[0] == 0) & (n_shared == 5)

    delta, n_shared = calculate_distance(a, b, offset_dict)
    assert (delta[0] == -1.0) & (n_shared == 4)

test_calcualte_distance()


# +
#export 

def calib_table(table: pd.DataFrame, delta: pd.Series, offset_dict: dict):
    """
    Apply offset to a table. Different operations for offsets exist.
    Offsets will be saved with a '_calib'-suffix. If this does not already exist,
    it will be created.

    Args:
        table_1 (pd.DataFrame): Dataframe with data.
        delta (pd.Series): Series cotaining the offset.
        offset_dict (dict): Dictionary with column names and how the distance should be calculated.

    Raises:
        NotImplementedError: If the type of vonversion is not implemented.
    """
    for col in offset_dict:

        if (col not in table.columns) and (col+'_apex' in table.columns):
            col_ = col+'_apex'
        else:
            col_ = col

        if offset_dict[col] == 'absolute':
            table[col+'_calib'] =  table[col_]-delta[col]
        elif offset_dict[col] == 'relative':
            table[col+'_calib'] = (1-delta[col_])*table[col]
        else:
            raise NotImplementedError(offset_dict[col])


# +
#hide
def test_calib_table():
    a = pd.DataFrame({'precursor':['A','B','C','D','E'], 'mass':[10, 20, 30, 40, 50]}).set_index('precursor')
    delta = pd.Series({'mass':[1]})
    offset_dict = {'mass':'absolute'}
    calib_table(a, delta, offset_dict)

    assert 'mass_calib' in a
    assert np.allclose(a['mass_calib'], a['mass']-1)

test_calib_table()

# +
#export 
import logging
from sklearn.linear_model import LinearRegression

def align(deltas: pd.DataFrame, filenames: list, weights:np.ndarray=None) -> np.ndarray:
    """Align multiple datasets.
    This function creates a matrix to represent the shifts from each dataset to another.
    This effectively is an overdetermined equation system and is solved with a linear regression.

    Args:
        deltas (pd.DataFrame): Distances from each dataset to another.
        filenames (list): The filenames of the datasts that were compared.
        weights (np.ndarray, optional): Distances can be weighted by their number of shared elements. Defaults to None.

    Returns:
        np.ndarray: alignment values.
    """
    matrix = []

    for i in range(len(deltas)):
        start, end = deltas.index[i]

        start_idx = filenames.index(start)
        end_idx = filenames.index(end)

        lines = np.zeros(len(filenames)-1)
        lines[start_idx:end_idx] = 1
        matrix.append(lines)

    # Remove nan values
    not_nan = ~deltas.isnull().any(axis=1)
    matrix = np.array(matrix)
    matrix = matrix[not_nan]
    deltas_ = deltas[not_nan]

    if len(deltas) < matrix.shape[1]:
        logging.info('Low overlap between datasets detected. Alignment may fail.')

    if weights is not None:
        reg = LinearRegression(fit_intercept=False).fit(matrix, deltas_.values, sample_weight = weights[not_nan])
        score= reg.score(matrix, deltas_.values)
    else:
        reg = LinearRegression(fit_intercept=False).fit(matrix, deltas_.values)
        score= reg.score(matrix, deltas_.values)

    logging.info(f"Regression score is {score}")

    x = reg.predict(np.eye(len(filenames)-1))

    return x


# +
#hide
def test_align():
    deltas = pd.DataFrame({'filename':[('A', 'B'), ('A', 'C'), ('B', 'C')],'mass_offset': [1, -1, -2] }).set_index('filename')
    filenames = ['A','B','C']

    assert np.allclose(align(deltas, filenames), np.array([[1],[-2]]))

test_align()

# +
#export 
import alphapept.io
import os
from typing import Callable

def calculate_deltas(combos: list, calib:bool = False, callback:Callable=None) -> (pd.DataFrame, np.ndarray, dict):

    """Wrapper function to calculate the distances of multiple files.
    
    In here, we define the offset_dict to make a relative comparison for mz and mobility and absolute for rt.
    
    TODO: This function could be speed-up by parallelization

    Args:
        combos (list): A list containing tuples of filenames that should be compared.
        calib (bool): Boolean flag to indicate distance should be calculated on calibrated data.
        callback (Callable): A callback function to track progress.

    Returns:
        pd.DataFrame: Dataframe containing the deltas of the files
        np.ndarray: Numpy array containing the weights of each comparison (i.e. number of shared elements)
        dict: Offset dictionary whicch was used for comparing. 

    """
    
    offset_dict = {}
    deltas = pd.DataFrame()
    weights = []

    for i, combo in enumerate(combos):
        file1 = os.path.splitext(combo[0])[0] + '.ms_data.hdf'
        file2 = os.path.splitext(combo[1])[0] + '.ms_data.hdf'
        df_1 = alphapept.io.MS_Data_File(file1).read(dataset_name="peptide_fdr").set_index('precursor')
        df_2 = alphapept.io.MS_Data_File(file2).read(dataset_name="peptide_fdr").set_index('precursor')

        if not offset_dict:
            offset_dict = {'mz':'relative', 'rt':'absolute'}
            if 'mobility' in df_1.columns:
                logging.info("Also using mobility for calibration.")
                offset_dict['mobility'] = 'relative'
            cols = list(offset_dict.keys())

        if len(deltas) == 0:
             deltas = pd.DataFrame(columns = cols)

        dists, weight = calculate_distance(df_1, df_2, offset_dict, calib = calib)
        deltas = deltas.append(pd.DataFrame([dists], columns = cols, index=[combo]))

        weights.append(weight)

        if callback:
            callback((i+1)/len(combos))

    return deltas, np.array(weights), offset_dict


# -

import alphapept.io


# +
#hide
def test_calculate_deltas():
    filename = '../testfiles/test.raw'
    combos = [(filename, filename)]
    deltas, weights, offset_dict = calculate_deltas(combos)
    assert np.allclose(deltas.values, np.array((0,0)))

test_calculate_deltas()

# +
#export
import pandas as pd
from itertools import combinations
import numpy as np
import os
import functools

#There is no unit test for align_files and align_datasets as they are wrappers and should be covered by the quick_test
def align_files(filenames: list, alignment: pd.DataFrame, offset_dict: dict):
    """
    Wrapper function that aligns a list of files.
    
    Args:
        filenames (list): A list with raw file names.
        alignment (pd.DataFrame): A pandas dataframe containing the alignment information.
        offset_dict (dict): Dictionary with column names and how the distance should be calculated.
    """
    for idx, filename in enumerate(filenames):

        file = os.path.splitext(filename)[0] + '.ms_data.hdf'

        for column in ['peptide_fdr', 'feature_table']:
            df = alphapept.io.MS_Data_File(file).read(dataset_name=column)
            calib_table(df, alignment.iloc[idx], offset_dict)
            logging.info(f"Saving {file} - {column}.")
            ms_file = alphapept.io.MS_Data_File(file, is_overwritable=True)

            ms_file.write(df, dataset_name=column)


def align_datasets(settings:dict, callback:callable=None):
    """
    Wrapper function that aligns all experimental files specified a settings file.
    
    Args:
        settings (dict): A list with raw file names.
        callback (Callable): Callback function to indicate progress.
    """
    filenames = settings['experiment']['file_paths']

    if callback:
        def progress_wrapper(current, step, n_steps):
            callback((step/n_steps)+(current/n_steps))

        cb = functools.partial(progress_wrapper, 0, 2)
    else:
        cb = None

    if len(filenames) > 1:
        combos = list(combinations(filenames, 2))

        deltas, weights, offset_dict = calculate_deltas(combos, callback=cb)

        cols = list(offset_dict.keys())

        before_sum = deltas.abs().sum().to_dict()
        before_mean = deltas.abs().mean().to_dict()

        logging.info(f'Total deviation before calibration {before_sum}')
        logging.info(f'Mean deviation before calibration {before_mean}')

        logging.info(f'Solving equation system')

        alignment = pd.DataFrame(align(deltas, filenames, weights), columns = cols)
        alignment = pd.concat([pd.DataFrame(np.zeros((1, alignment.shape[1])), columns= cols), alignment])
        alignment -= alignment.mean()

        logging.info(f'Solving equation system complete.')

        logging.info(f'Applying offset')

        align_files(filenames, -alignment, offset_dict)

        if cb:
            cb = functools.partial(progress_wrapper, 1, 2)

        deltas, weights, offset_dict = calculate_deltas(combos, calib=True, callback=cb)

        after_sum = deltas.abs().sum().to_dict()
        after_mean = deltas.abs().mean().to_dict()

        logging.info(f'Total deviation after calibration {after_sum}')
        logging.info(f'Mean deviation after calibration {after_mean}')

        change_sum = {k:v/before_sum[k] for k,v in after_sum.items()}
        change_mean = {k:v/before_mean[k] for k,v in after_mean.items()}

        logging.info(f'Change (after/before) total deviation {change_sum}')
        logging.info(f'Change (after/before) mean deviation {change_mean}')

    else:
        logging.info('Only 1 dataset present. Skipping alignment.')


# -

# ## Matching 
#
# Transfer MS2 identifications to similar MS1 features.
#
# For "match-between-runs" we start with aligning datasets. To create a reference we use for matching, we combine all datasets of a matching group. When using the default settings, the matching group consists of all files. We then group the dataset by precursor and calculate it's average properties (rt, mz, mobility). By combining several files we further are able to calculate a standard deviation. This allows us to know where and with what deviation we would expect an MS1 feature and have the corresponding identification. This is our matching reference. In the matching step, we go through each dataset individually and check if there are precursors in the reference that were not identified in this dataset. We then perform a nearest-neighbor lookup to find if any MS1 features exist that are in close proximity to the reference. The distance metric we use is normed by the median standard of the deviation. Lastly we assess the confidence in a transfered identifcation by using the [`Mahalanobis`](https://stats.stackexchange.com/questions/331283/how-to-calculate-the-probability-of-a-data-point-belonging-to-a-multivariate-nor) distance.

#export
from scipy import stats
def get_probability(df: pd.DataFrame, ref: pd.DataFrame, sigma:pd.DataFrame, index:int)-> float:
    """Probablity estimate of a transfered identification using the Mahalanobis distance.
    
    The function calculates the probability that a feature is a reference feature.
    The reference features containing std deviations so that a probability can be estimated.
    
    It is required that the data frames are matched, meaning that the first entry in df matches to the first entry in ref.
    
    Args:
        df (pd.DataFrame): Dataset containing transferered features
        ref (pd.DataFrame): Dataset containing reference features
        sigma (pd.DataFrame): Dataset containing the standard deviations of the reference features
        index (int): Index to the datframes that should be compared

    Returns:
        float: Mahalanobis distance
    """

    sigma = sigma.iloc[index].values
    sigma = sigma*np.eye(len(sigma))

    mu = ref.iloc[index].values
    x = df.iloc[index].values

    try:
        m_dist_x = np.dot((x-mu).transpose(), np.linalg.inv(sigma))
        m_dist_x = np.dot(m_dist_x, (x-mu))
        _ = stats.chi2.cdf(m_dist_x, len(mu))
    except Exception as e:
        _ = np.nan

    return _


# +
#Example usage

a = pd.DataFrame({'mass':[100,200,300],'rt':[1,2,3]})
b = pd.DataFrame({'mass':[100,200,302],'rt':[1,2.5,3]})
std = pd.DataFrame({'mass':[0.1,0.1,0.1],'rt':[1,1,1]})

print(f"First element: (ideal match): {get_probability(a, b, std, 0):.2f}")
print(f"Second element: (rt slightly off): {get_probability(a, b, std, 1):.2f}")
print(f"Third element: (mass completely off): {get_probability(a, b, std, 2):.2f}")


# +
#hide
def test_get_probability():
    a = pd.DataFrame({'mass':[100,200,300],'rt':[1,2,3]})
    b = pd.DataFrame({'mass':[100,200,302],'rt':[1,2.5,3]})
    std = pd.DataFrame({'mass':[0.1,0.1,0.1],'rt':[1,1,1]})

    assert np.allclose(get_probability(a, b, std, 0), 0.00)
    assert np.allclose(get_probability(a, b, std, 2), 1)
    
test_get_probability()

# +
#export 
from sklearn.neighbors import KDTree
from alphapept.utils import assemble_df

# This function is a wrapper function and has currently has no unit test 
# The function will be revised when implementing issue #255: https://github.com/MannLabs/alphapept/issues/255
def match_datasets(settings:dict, callback:Callable = None):
    """Match datasets: Wrapper function to match datasets based on a settings file.
    
    Args:
        settings (dict): Dictionary containg specifications of the run
        callback (Callable): Callback function to indicate progress.
    """

    if len(settings['experiment']['file_paths']) > 2:
        xx = alphapept.utils.assemble_df(settings, field='peptide_fdr')

        base_col = ['precursor']
        alignment_cols = ['mz_calib','rt_calib']
        extra_cols = ['score','decoy','target']

        if 'mobility' in xx.columns:
            alignment_cols += ['mobility_calib']
            use_mobility = True
        else:
            use_mobility = False

        grouped = xx[base_col + alignment_cols + extra_cols].groupby('precursor').mean()
        std_ = xx[base_col + alignment_cols].groupby('precursor').std()

        grouped[[_+'_std' for _ in alignment_cols]] = std_

        std_range = np.nanmedian(std_.values, axis=0)

        match_p_min = settings['matching']['match_p_min']
        match_d_min = settings['matching']['match_d_min']

        filenames = settings['experiment']['file_paths']

        lookup_dict = xx.set_index('precursor')[['sequence']].to_dict()

        for idx, filename in enumerate(filenames):
            file = os.path.splitext(filename)[0] + '.ms_data.hdf'

            df = alphapept.io.MS_Data_File(file).read(dataset_name='peptide_fdr')
            features = alphapept.io.MS_Data_File(file).read(dataset_name='feature_table')
            features['feature_idx'] = features.index

            matching_set = set(grouped.index) - set(df['precursor'])
            logging.info(f'Trying to match file {file} with database of {len(matching_set):,} unidentified candidates')

            mz_range = std_range[0]
            rt_range = std_range[1]

            tree_points = features[alignment_cols].values
            tree_points[:,0] = tree_points[:,0]/mz_range
            tree_points[:,1] = tree_points[:,1]/rt_range

            query_points = grouped.loc[matching_set][alignment_cols].values
            query_points[:,0] = query_points[:,0]/mz_range
            query_points[:,1] = query_points[:,1]/rt_range

            if use_mobility:
                logging.info("Using mobility")
                i_range = std_range[2]

                tree_points[:,2] = tree_points[:,2]/i_range
                query_points[:,2] = query_points[:,2]/i_range

            matching_tree = KDTree(tree_points, metric="minkowski")

            dist, idx = matching_tree.query(query_points, k=1)

            matched = features.iloc[idx[:,0]]

            for _ in extra_cols:
                matched[_] = grouped.loc[matching_set, _].values

            to_keep = dist < match_d_min

            matched = matched[to_keep]

            ref = grouped.loc[matching_set][alignment_cols][to_keep]
            sigma = std_.loc[matching_set][to_keep]

            logging.info(f'{len(matched):,} possible features for matching based on distance of {match_d_min}')

            matched['matching_p'] = [get_probability(matched[alignment_cols], ref, sigma, i) for i in range(len(matched))]
            matched['precursor'] = grouped.loc[matching_set][to_keep].index.values

            matched = matched[matched['matching_p']< match_p_min]

            logging.info(f'{len(matched):,} possible features for matching based on probability of {match_p_min}')

            matched['type'] = 'matched'

            for _ in lookup_dict.keys():
                matched[_] = [lookup_dict[_][x] for x in matched['precursor']]

            df['type'] = 'msms'
            df['matching_p'] = np.nan

            shared_columns = set(matched.columns).intersection(set(df.columns))

            df_ = pd.concat([df, matched[shared_columns]], ignore_index=True)

            logging.info(f"Saving {file} - peptide_fdr.")
            ms_file = alphapept.io.MS_Data_File(file, is_overwritable=True)

            ms_file.write(df_, dataset_name='peptide_fdr')
    else:
        logging.info('Less than 3 datasets present. Skipping matching.')


# -

#hide
from nbdev.export import *
notebook2script()

StarcoderdataJupyterScriptsDedupFiltered