"""Zooming into high resolution mass spectrometry data"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/05_exploring.ipynb.

# %% auto #0
__all__ = ['read_raw', 'read_mzml', 'histogram', 'get_time_averaged_centroids', 'interactive_plot']

# %% ../notebooks/05_exploring.ipynb #945d18a4-62f9-4445-964a-14b5c1aa966b
import pyRawTools as rwt 
import dask.array as da 
import scipy.signal as ssg 
from scipy.ndimage import gaussian_filter 
import numpy as np 

import datashader as ds
import datashader.transfer_functions as tf
from datashader.mpl_ext import dsshow 
import matplotlib.cm as cm 
import matplotlib.pyplot as plt 

from functools import partial  

# %% ../notebooks/05_exploring.ipynb #774de0d5-dca0-4123-83a7-d3de3c07aad7
def read_raw(raw_file): 
    '''Read `raw_file` into positive and negative mode data frames. '''

    loader = rwt.MSLoader()
    df = loader.load(raw_file) 
    df.columns = ['RT', 'mz', 'inty']
    df_pos = df.iloc[::2]
    df_neg = df.iloc[1::2] 

    return df_pos, df_neg
    

def read_mzml(mzml_file): 
    '''Read `mzml_file`. 
    
    Returns positive and negative mode dataframes `df_pos` and `df_min`. '''

    # load example file
    exp = oms.MSExperiment()
    oms.MzMLFile().load(mzml_file, exp) 
    df_long = exp.get_df(long=True) 

    # split +/- 
    df_pos = df_long.iloc[::2] # positive?
    df_neg = df_long.iloc[1::2] # negative? 

    return df_pos, df_neg 


def histogram(df):
    '''Create intensity weighed histogram. '''

    # create dask arrays 
    x = da.from_array(df['mz'])
    y = da.from_array(df['inty']) 
    
    # make dask histograms 
    #mz_hist, bins = da.histogram(x, bins=10000000, range=[0, 1000]) # simply number of mass points in bin 
    mz_hist_w, bins = da.histogram(x, weights=y, bins=10000000, range=[0, 1000]) # weighed with intensities 
    
    # bin centers for plotting 
    bin_delta = bins[1] - bins[0]
    bin_centers = bins[0:-1] + 0.5 * bin_delta 

    mz_hist_w_xy = np.c_[bin_centers, mz_hist_w]

    return mz_hist_w_xy

def get_time_averaged_centroids(mz_hist_w_xy):
    '''Get peaks (centroids) from histogram. '''
    
    bin_centers, mz_hist_w = mz_hist_w_xy.T 
    
    # gaussian smoothing of weighed histogram 
    mz_hist_w_gauss = gaussian_filter(mz_hist_w, sigma=2)
    
    # find peaks 
    peak_idxs_w_gauss, props_dict_w = ssg.find_peaks(mz_hist_w_gauss, height=4000, distance=50) # bit ad hoc 
    peak_heights_w_gauss = props_dict_w['peak_heights'] 

    mz_x = bin_centers[peak_idxs_w_gauss] 
    mz_y = peak_heights_w_gauss 

    mz_centroids = np.c_[mz_x, mz_y]
    
    return mz_centroids  

def interactive_plot(df, mz_hist, mz_centroids, title=None): 
    '''Create interactive plot for dataframe `df`. '''

    mz_hist_x, mz_hist_y = mz_hist.T
    mz_centroids_x, mz_centroids_y = mz_centroids.T
    
    
    fig, [ax1, ax2] = plt.subplots(nrows=2, sharex=True, figsize=[8, 6], gridspec_kw={'height_ratios':[4, 1]}, squeeze=True) 
    
    # number histogram 
    #ax0.plot(bin_centers, mz_hist, c='r')
    #ymin, ymax = ax0.get_ylim()
    #ax0.vlines(bin_centers[peak_idxs_w_gauss], ymin=ymin, ymax=ymax, color=[0.9, 0.9, 0.9], zorder=-10) 
    #ax0.set_ylabel('bin count')
    #ax0.set_title('Ref0443_casein_asap01.mzML (+)')
    mz_min = df['mz'].min()
    
    
    # datashader beast (the most complicated single plot command I have ever made)
    glyph = ds.Point('mz', 'RT') 
    dsshow(df, glyph, aggregator=ds.sum('inty'), ax=ax1, x_range=[100, 1000], plot_height=500, 
           plot_width=1000, aspect='auto', cmap=cm.Reds, norm='log',  
           shade_hook=partial(tf.dynspread, threshold=0.7, max_px=7, how='add')) 
    ymin, ymax = ax1.get_ylim()
    ax1.vlines(mz_centroids_x, ymin=ymin, ymax=ymax, color=[0.9, 0.9, 0.9], zorder=-10)
    ax1.set_ylabel('RT')
    if title != None: 
        ax1.set_title(title)
    
    # intensity weighed histogram 
    ax2.plot(mz_hist_x, mz_hist_y, color='r')
    ymin, ymax = ax2.get_ylim()
    ax2.vlines(mz_centroids_x, ymin=ymin, ymax=ymax, color=[0.9, 0.9, 0.9], zorder=-10) 
    ax2.set_ylabel('inty')
    ax2.set_xlabel('m/z')
    ax2.set_xlim([mz_min, 1000])
    
    plt.subplots_adjust(hspace=0.02)
