"""But keeping time"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/07_centroiding-again.ipynb.

# %% auto #0
__all__ = ['centralize']

# %% ../notebooks/07_centroiding-again.ipynb #4ec6e619-e322-4bd2-956d-ef67977fc6c3
from dask import array as da 

# %% ../notebooks/07_centroiding-again.ipynb #85d7ce9a-9b3e-443a-a809-9022c059e0a7
def centralize(df, mz_centroids, precision=0.002):
    '''Shift mass points to centroid positions. 
    
    Keep retention time dimension.'''
    
    # extract numpy arrays from dataframe and centroids array 
    mz = df['mz'].values 
    mzc = mz_centroids[:,0] 
    
    # for efficient computation convert to dask arrays 
    mz_arr = da.from_array(mz)
    mzc_arr = da.from_array(mzc)
    
    # create distances table 
    diff_arr = ((mz_arr[:, None] - mzc_arr[None, :])**2)**0.5
    
    # compute nearest distance for each mass point 
    distances = diff_arr.min(axis=1).compute()
    
    # within precision 
    is_nearby = distances < precision
    
    # locate corresponding centroid mass indexes 
    nearest_idxs = list(diff_arr.argmin(axis=1).compute())
    
    # replace mz column in dataframe with centroid mz 
    df_centr = df.copy()
    df_centr['mz'] = mzc[nearest_idxs]
    
    # discard outliers
    df_centr = df_centr[is_nearby]

    return df_centr
