"""Dealing with mixed signals"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/60_peakmaps.ipynb.

# %% auto 0
__all__ = ['get_continuum', 'get_gaussians', 'fit_spectrum_peaks', 'compute_nmf_peakmaps', 'multi_plot',
           'compute_nmf_element_maps']

# %% ../notebooks/60_peakmaps.ipynb 30
from dask import delayed 
from dask.diagnostics import ProgressBar 
import maxrf4u 
import numpy as np 
import scipy.interpolate as sip 
import scipy.signal as ssg 
import matplotlib.pyplot as plt 
from IPython.display import SVG 

import sklearn.decomposition as skd 

# %% ../notebooks/60_peakmaps.ipynb 31
def get_continuum(datastack_file): 
    '''Compute continuum baseline from sum spectrum.
    
    Uses rolling ball filter to remove peaks from sum spectrum. 
    
    Returns: y_continuum
    '''
    
    ds = maxrf4u.DataStack(datastack_file) 
    x_keVs = ds.read('maxrf_energies')
    y_sumspectrum = ds.read('maxrf_sumspectrum') 
    
    # rolling ball filter
    y_continuum = maxrf4u.hotmax._baseline(x_keVs, y_sumspectrum) # radius=200 
    # repeat with smaller ball to cut sharp corners of Compton 
    y_continuum =  maxrf4u.hotmax._baseline(x_keVs, y_continuum, radius=40)
    
    return y_continuum 


def _gaussian(x, x0, sigma): 
    '''Normal distribution around `x0` with standard deviation `sigma`.'''
    
    y = np.exp(-0.5 * ((x - x0) / sigma)**2)
    
    return y 


def _fit_gaussian(x, y, peak_idx, rel_height=0.2, baseline=None): 
    '''Fit single gaussian distribution at `rel_height`. 
    
    Returns: `y_gauss`, `baseline`
    '''
    
    # just a single peak 
    [widths], [width_height], [left_ips], [right_ips] = ssg.peak_widths(y, [peak_idx], rel_height=rel_height)
    
    # create energy interpolation 
    keV_ipol = sip.interp1d(np.arange(len(x)), x)
     
    left_x = keV_ipol(left_ips)
    right_x = keV_ipol(right_ips)
    
    # pick smallest half width at rel_height 
    dx_left = left_x - x[peak_idx] 
    dx_right = right_x - x[peak_idx] 
    
    dx = sorted([dx_left, dx_right])[0] 
    
    # create baseline if None   
    if baseline is None:  
        baseline = np.zeros_like(y) 
        
    y_norm = (width_height - baseline[peak_idx]) / (y[peak_idx] - baseline[peak_idx]) 
    
    # assert that y_norm is positive 
    assert y_norm > 0, f"Can't fit Gaussian for peak below baseline for peak index {peak_idx}."   
        
    # calculate corresponding sigma 
    sigma = np.sqrt(-dx**2 / (2 * np.log(y_norm)))
    
    # calculate gaussian with baseline 
    y_gauss = (y[peak_idx] - baseline[peak_idx]) * _gaussian(x, x[peak_idx], sigma) + baseline
    
    return y_gauss, baseline 


def get_gaussians(datastack_file, tail_clip=0.05, norm=True): 
    '''Fit Gaussian peak profiles to hotmax peaks in hotmax atlas. 

    Returns: y_gauss_list
    '''

    # read stuff from datastack 
    ds = maxrf4u.DataStack(datastack_file) 

    x_keVs = ds.read('maxrf_energies') 
    hotmax_spectra = ds.read('hotmax_spectra')
    hotmax_baselines = ds.read('hotmax_baselines') 
    peak_idxs_list = ds.read_list('hotmax_peak_idxs_list')

    # get slices by fitting gaussian to corresponding hotmax spectrum and baseline 
    y_gauss_list = []

    # Create tail clipped gaussians for ragged peak indexes list 
    for n, peak_idxs in enumerate(peak_idxs_list):  

        y_hot = hotmax_spectra[n]
        baseline = hotmax_baselines[n] 

        # iterate over (possibly multiple) peaks in each hotmax spectrum 
        for idx in peak_idxs: 

            y_gauss_fit, baseline = _fit_gaussian(x_keVs, y_hot, idx, baseline=baseline) 
            y_gauss_flat =  y_gauss_fit - baseline 

            if norm is True: 
                y_gauss_flat = y_gauss_flat / y_gauss_flat.max()

            y_gauss_list.append(y_gauss_flat)
            
    return y_gauss_list 


def fit_spectrum_peaks(y, datastack_file): 
    '''Fit Gaussian peak shapes to single spectrum `y`.
    
    Uses NMF with fixed Gaussian component vectors. 
    
    Returns: weighed_components_list 
    '''

    X = y.reshape([1, -1]) # make single spectro into two dimenional array 
    

    y_continuum = get_continuum(datastack_file)
    y_gauss_list = get_gaussians(datastack_file, norm=False)

    # create component vectors 
    H = np.array([y_gauss for y_gauss in y_gauss_list]).astype(np.float32)
    # add continuum as component
    H = np.r_[H, y_continuum[None,:].astype(np.float32)]
    
    n_components, n_channels = H.shape

    W, H, n_iter = skd.non_negative_factorization(X, H=H, n_components=n_components, update_H=False)
    
    # not neede here, but just a reminder  
    # X_nmf = W @ H 
    # y_nmf = X_nmf.flatten() 

    # transpose weights for easy multiplication  
    w =  W.T 
    
    return w, H 
        

    
def _add_hotlines_ticklabels(datastack_file, ax, vlines=True, clip_vline=True): 
    '''Utility function. Adds hotlines and tick labels to plot `ax`. 
    
    '''
    # read hotlines data 
    ds = maxrf4u.DataStack(datastack_file) 
    x_keVs = ds.read('maxrf_energies')
    peak_idxs = ds.read('hotmax_peak_idxs_flat') 
    
    secax = ax.secondary_xaxis('top')

    secax.set_xticks(x_keVs[peak_idxs]) 
    
    ymin, ymax = ax.get_ylim() 
    
    # include vlines in plot 
    if vlines: 
        # clip vlines at y=0
        if clip_vline: 
            ax.vlines(x_keVs[peak_idxs], 0, 1.2*ymax, color='r', alpha=0.2, zorder=9-30)
            ax.set_ylim(ymin, ymax)

        # do not clip vlines 
        else: 
            for x in x_keVs[peak_idxs]:
                ax.axvline(x, color='r', alpha=0.2, zorder=9-30)

    secax.set_xticks(x_keVs[peak_idxs])
    secax.set_xticklabels(range(len(peak_idxs)), fontsize=6, color='r') 
    secax.tick_params(color=[1, 0.5, 0.5], pad=0) 
    
    return ax 


def _get_spectral_slices(cube): 
    '''Low level bookkeeping to compute slice indices for spectral chunks from dask array `cube`. 
       
    Returns: row_slices, column_slices  
    '''
    
    h, w, d = cube.shape 

    row_chunks, column_chunks, spectral_chunks = cube.chunks

    row_idxs = [0] 
    for c in row_chunks: 
        row_idxs.append(row_idxs[-1] + c)

    row_slices = [[row_idxs[i-1], row_idxs[i]] for i in range(1, len(row_idxs))]

    # columns 
    column_idxs = [0] 
    for c in column_chunks: 
        column_idxs.append(column_idxs[-1] + c)

    column_slices = [[column_idxs[i-1], column_idxs[i]] for i in range(1, len(column_idxs))] 

    return row_slices, column_slices 



def compute_nmf_peakmaps(datastack_file): 
    '''Use nonnegative matrix factorization to compute peak maps spectral data in `datastack_file`. 
    
    Asks to write result to datastack file. 
    
    Returns: peak_maps 
    '''
    
    ds = maxrf4u.DataStack(datastack_file)
    cube = ds.read('maxrf_cube', compute=False)
    h, w, d = cube.shape 
    
    
    # prepare gaussian component vectors 
    y_continuum = get_continuum(datastack_file)
    y_gauss_list = get_gaussians(datastack_file)

    H = np.array([y_gauss for y_gauss in y_gauss_list]).astype(np.float32)
    # add continuum as component
    H = np.r_[H, y_continuum[None,:].astype(np.float32)]

    n_components, n_channels = H.shape

    #initialize peak map stack  
    peak_map_stack = np.zeros([h, w, n_components], dtype=np.float32)

    # prepare chunk bookkeeping 
    row_slices, column_slices = _get_spectral_slices(cube)

    n_chunks = len(row_slices) * len(column_slices)

    n = 1 
    print('Please wait a few minutes while preparing your peak maps. ')
    for ri, rj in row_slices: 

        for ci, cj in column_slices: 
            print(f'Computing factorization for spectral chunk {n}/{n_chunks}...', end='\r')

            spectral_chunk = cube[ri:rj, ci:cj, :].compute()
            h, w, d = spectral_chunk.shape
            X = spectral_chunk.reshape([h*w, d])

            W, H, n_iter = skd.non_negative_factorization(X, H=H, n_components=n_components, update_H=False) 

            W_img = W.reshape([h, w, -1]) 

            peak_map_stack[ri:rj, ci:cj] = W_img 
            
            n = n + 1

    peak_maps = peak_map_stack.transpose([2, 0, 1])
    
    # user input           
    write = input('Write peak maps factorization to datastack file [y/n]? ')
         
    if write == 'y': 
        maxrf4u.append(peak_maps, 'nmf_peakmaps', datastack_file)
        maxrf4u.append(H, 'nmf_gausscomponents', datastack_file) 
        
        print(f'\nSaved NMF peak maps data to: {datastack_file}')
    
    return peak_maps 


def multi_plot(*images, hot_pixel=None, titles=None, roi_list=None, axis_off=False,
               sharex=True, sharey=True, vmin=None, vmax=None, cmap='viridis',
               fontsize='medium', zoom_xyc=None, zoom_half_wh=[100, 100], svg=False):
    '''Inspect multiple images simultaneously...
    
    Fold along multiple rows if n > 4'''
    
    if svg:
        plt.ioff()
    
    nrows_max = 4
    n_img = len(images)
    nrows = (n_img // nrows_max) # completely filled rows
    rest = n_img % nrows_max
    if rest != 0:
        nrows = nrows + 1
    if n_img <= nrows_max:
        ncols = n_img
    else:
        ncols = nrows_max
    figsize = [9, 5 + 1.3 * (nrows -1)]
    fig, axs = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize, squeeze=False, sharex=sharex, sharey=sharey)
    for i, img in enumerate(images):
        axs.flatten()[i].imshow(img, cmap=cmap, vmin=vmin, vmax=vmax)
        if hot_pixel is not None:
            hot_yi, hot_xi, hot_zi = hot_pixel
            axs.flatten()[i].scatter(hot_xi, hot_yi, color='r')
        if zoom_xyc is not None:
            xc, yc = zoom_xyc
            w_ha, h_ha = zoom_half_wh
            axs.flatten()[i].set_xlim(xc - w_ha, xc + w_ha)
            axs.flatten()[i].set_ylim(yc + h_ha, yc - h_ha)
        if roi_list is not None:
            add_roi_patches(axs.flatten()[i], roi_list)
    if titles is not None:
        for i, t in enumerate(titles):
            axs.flatten()[i].set_title(t, fontsize=fontsize)
    # seems to have no effect:
    for i in range(n_img, nrows * ncols):
        axs.flatten()[i].axis('off')
    # therefore trying this:
    if axis_off:
        axs_flat = axs.flatten()
        for ax in axs_flat:
            ax.set_axis_off()
    fig.subplots_adjust(hspace=0.1, wspace=0.03)
    plt.tight_layout()
    
    
    # larger image in documentation 
    if svg: 
        plt.savefig('plot.svg')
        fig.clear()
        plt.ion()
        return SVG('plot.svg')
    
    return fig, axs


def compute_nmf_element_maps(datastack_file, elements_unsorted, excitation_energy_keV=25): 
    '''Compute element maps for `elements` list. 
    
    Based on NMF factorization. Requires previously computed peak maps stored in `datastack_file`. 
    Sorts elements according to strongest (a.k.a. alpha) peak. Ask user to save results. 
    
    Returns: elements, element_maps
    '''

    ds = maxrf4u.DataStack(datastack_file)
    
    x_keVs = ds.read('maxrf_energies')

    nmf_gausscomponents_plus_continuum = ds.read('nmf_gausscomponents')
    nmf_gausscomponents = nmf_gausscomponents_plus_continuum[0:-1]

    # the data X that we want to factorize first is the element_spectra
    
    print('Preparing peaks to elements transformation matrix...')
    
    elements, element_spectra = maxrf4u.get_element_spectra(elements_unsorted, x_keVs, 
                                                            excitation_energy_keV)
    
    X = element_spectra.astype(np.float32) 
    H = nmf_gausscomponents 
    n_components, n_channels = H.shape 
    W_elem, H, n_iter = skd.non_negative_factorization(X, H=H, n_components=n_components, update_H=False) 

    # Ok, this might be slightly confusing, but now this weights matrix becomes the new components matrix. 

    peak_maps = ds.read('nmf_peakmaps')

    # Need to transpose first, then temporarily flatten.

    peak_maps_stack = peak_maps.transpose([1, 2, 0])

    h, w, d = peak_maps_stack.shape
    
    # this is actual the NMF transformation of peak maps into element maps  
    
    X = peak_maps_stack.reshape([h * w, d])  # flatten 
    X = X[:,0:-1] # get rid of continuum layer 

    H = W_elem 
    n_components, n_peaks = H.shape 
    
    print(f'Computing element maps for {elements}...')

    W2, H, n_iter = skd.non_negative_factorization(X, H=H, n_components=n_components, update_H=False)
    
    print('Ready computing!\n')

    # Ok, unflatten and transpose again. 

    element_maps = W2.reshape([h, w, -1]).transpose([2, 0, 1])
    
    # user input           
    write = input('Write NMF element maps and factorization data to datastack file [y/n]? ')
         
    if write == 'y': 
        maxrf4u.append(element_maps, 'nmf_elementmaps', datastack_file)
        maxrf4u.append(np.array(elements), 'nmf_elements', datastack_file) 
        maxrf4u.append(W_elem, 'nmf_peaks2elements_matrix', datastack_file)
        
        print(f'\nSaved NMF element maps and factorization data to: {datastack_file}')

    return elements, element_maps 
