"""Matching peak patterns"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/50_peak-pattern-puzzle.ipynb.

# %% auto 0
__all__ = ['all_elements', 'eoi', 'get_patterns', 'colorize', 'plot_ptrn', 'plot_patterns', 'plot_puzzle',
           'get_instrument_pattern', 'add_hotlines']

# %% ../notebooks/50_peak-pattern-puzzle.ipynb 50
import maxrf4u  

from . import HotmaxAtlas   
from .peakmaps import _add_hotlines_ticklabels 

from .xphysics import get_element_spectrum 


import re 
import matplotlib.pyplot as plt 
import matplotlib.cm as cm 
import mendeleev
import moseley as mos
import numpy as np 

# %% ../notebooks/50_peak-pattern-puzzle.ipynb 51
# all elements 
all_elements = ['#H', '#He', '#Li', '#Be', '#B', '#C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 
                'Si', 'P', 'S', 'Cl', '#Ar', 'K', 'Ca', '#Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 
                'Co', 'Ni', 'Cu', 'Zn', '#Ga', '#Ge', 'As', '#Se', 'Br', '#Kr', '#Rb', 'Sr', 
                '#Y', '#Zr', '#Nb', '#Mo', '#Tc', '#Ru', 'Rh', '#Pd', 'Ag', 'Cd', '#In', 'Sn', 
                '#Sb', '#Te', 'I', '#Xe', '#Cs', 'Ba', '#La', '#Hf', '#Ta', '#W', '#Re', '#Os', 
                '#Ir', '#Pt', '#Au', 'Hg', '#Tl', 'Pb', '#Bi', '#Po', '#At', '#Rn', '#Fr', '#Ra', 
                '#Ac', '#Rf', '#Db', '#Sg', '#Bh', '#Hs', '#Mt', '#Ds', '#Rg', '#Cn', '#Nh', 
                '#Fl', '#Mc', '#Lv', '#Ts', '#Og']

# elements of interest (to me) 
eoi = [e for e in all_elements if not '#' in e] 


def get_patterns(elements=None, tube_keV=30, color_select=None): 
    '''Returns pattern dict list for `elements` list, sorted according to alpha peak energy. 
    
    By default all elements are colored with a standard color. To colorize only specific elements, 
    provide a list of elements `color_select`.
    '''
     
    ptrn_dict_list = [] 
    
    if elements is None: 
        elements = eoi 
    
    for elem in elements: 

        spectrum = mos.XFluo(elem, tube_keV=tube_keV, min_prom=0.01)

        peak_order = np.argsort(spectrum.peak_intensities)[::-1]

        alpha_keV = spectrum.peak_energies[peak_order][0]
        peaks_x = spectrum.peak_energies[peak_order]
        peaks_y = spectrum.peak_intensities[peak_order]
        
        peaks_xy = np.c_[peaks_x, peaks_y] 
        
        alpha_escape_keV = alpha_keV - 1.74  # Silicon detector escape energy shift 
        
        if color_select is None:  
            # all elements are colored 
            color = colorize(elem) 
        else: 
            if elem in color_select: 
                # only selected elements are colored  
                color = colorize(elem) 
            else:
                # otherwise make black 
                color = [0, 0, 0]         
        
        
        name = mendeleev.element(elem).name
        
        ptrn_dict = {'elem': elem, 
                     'name': name, 
                     'peaks_xy': peaks_xy, 
                     'alpha_escape_keV': alpha_escape_keV, 
                     'color': color}
        
        ptrn_dict_list.append(ptrn_dict) 
        
    # sort according to alpha energy     
    alpha_keVs = [p['peaks_xy'][0, 0] for p in ptrn_dict_list]        
    indices = np.argsort(alpha_keVs) 
    
    ptrn_list = [ptrn_dict_list[i] for i in indices] 
        
    return ptrn_list  


def colorize(elem): 
    '''Pick fixed color from nice color map for elements of interest. '''
    
    # select elements of interest 
    all_elements = ['#H', '#He', '#Li', '#Be', '#B', '#C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 
                    'Si', 'P', 'S', 'Cl', '#Ar', 'K', 'Ca', '#Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 
                    'Co', 'Ni', 'Cu', 'Zn', '#Ga', '#Ge', 'As', '#Se', 'Br', '#Kr', '#Rb', 'Sr', 
                    '#Y', '#Zr', '#Nb', '#Mo', '#Tc', '#Ru', 'Rh', '#Pd', 'Ag', 'Cd', '#In', 'Sn', 
                    '#Sb', '#Te', 'I', '#Xe', '#Cs', 'Ba', '#La', '#Hf', '#Ta', '#W', '#Re', '#Os', 
                    '#Ir', '#Pt', '#Au', 'Hg', '#Tl', 'Pb', '#Bi', '#Po', '#At', '#Rn', '#Fr', '#Ra', 
                    '#Ac', '#Rf', '#Db', '#Sg', '#Bh', '#Hs', '#Mt', '#Ds', '#Rg', '#Cn', '#Nh', 
                    '#Fl', '#Mc', '#Lv', '#Ts', '#Og']

    eoi = [e for e in all_elements if not '#' in e] 
        
    # tab20x2 color map 
    tab20 = cm.tab20(np.arange(20))[:,0:3]
    colors = np.r_[tab20[::-1], tab20[::-1]**0.6]
    
    # this code logic is pretty ad hoc (needs improvement)
    # custom colors 
    special_colors = {'Pb': (0.4, 0.4, 0.4), 
                      'Cu': (0.1, 0.9, 0.3), 
                      'Fe': (0.7, 0.5, 0.1), 
                      'S': (1.0, 0.9, 0.1),
                      'Au': (1, 0.9, 0.1), 
                      'Br': (0.4, 0.3, 0)} 
    
    for e in special_colors.keys(): 
        if e in eoi: 
            colors[eoi.index(e)] = special_colors[e]
        
    if elem in eoi:  
            color = colors[eoi.index(elem)]  
            
    else: 
        color = [0, 0, 0]
    
    return color 
       

def plot_ptrn(elem, y, ax, color_select=None): 
    '''Low level plot element pattern at level `y` in axes `ax`.'''
    
    ptrn = get_patterns([elem], color_select=color_select)[0] 
    
    # would like to include the theoretical emission spectrum as well  
    # or simply do gaussian convolve? 
    x, elem_spectrum = get_element_spectrum(elem, excitation_energy_keV=30)
    
    peaks_x, peaks_y = ptrn['peaks_xy'].T  
    color = ptrn['color']
    
    alpha_escape_keV = ptrn['alpha_escape_keV'] 

    left_x = min(peaks_x)
    right_x = max(peaks_x)

    ones = np.ones_like(peaks_y) 
    
    # if below spectrum rescale height y according to available space 
    if y < 0: 
        ymin = ax.get_ylim()[0]
        y = -(y / 5) * ymin 

    ax.scatter(peaks_x, y * ones, s=15, color=color) 
    ax.plot([left_x, right_x], [y, y], color=color, alpha=0.3)
    
    
    ax.fill_between(x, y + 0.8*elem_spectrum, y2=y, alpha=0.2, color=color, 
                    where=elem_spectrum>0.01) # emission spectrum 
    ax.scatter(peaks_x[0], y, marker='s', s=40, color=color)
    
    ax.scatter(alpha_escape_keV, y, marker='|', s=15, color=color)
    
    ax.annotate(ptrn['elem'], [right_x, y], xytext=[5, -1], fontsize=8, color=color, 
                textcoords='offset points', ha='left', va='center') 
     
    return ptrn 

    
def plot_patterns(elements=None, instr_ptrn_from_datastack_file=None, color_select=None, ax=None): 
    '''Wrapper function to plot overview of `element` and instrument patterns in axes `ax`. 
    
    Returns: `ax` 
    '''
    
    if elements is None: 
        elements = eoi 
      
    ptrn_list = get_patterns(elements=elements, tube_keV=30, color_select=color_select) 
    
    # this should fix the alpha peak order tick label bug? 
    elements = [ptrn['elem'] for ptrn in ptrn_list]
    element_labels = [f'{mendeleev.element(e).name} ({e})' for e in elements] 
    
    n_ptrns = len(ptrn_list) 
    n_ticks = n_ptrns 
    offset = 0 
    ytick_labels = element_labels 
    
    if instr_ptrn_from_datastack_file is not None: 
        instrument_ptrn = get_instrument_pattern(instr_ptrn_from_datastack_file)
        
        n_ticks += 1 
        offset = 1 
        ytick_labels = ['INSTRUMENT'] + ytick_labels 
        
        instr_peaks = instrument_ptrn['instrument_peaks']
        instr_zeros = np.zeros_like(instr_peaks) 

    # now create plot    
    if ax is None:  
        fig_width = 9
        fig_height = 0.15 * n_ticks + 0.5 
        fig, ax = plt.subplots(figsize=[fig_width, fig_height])  
        ax.set_xlabel('energy (keV)') 
            
    # plot instrument pattern X X X 
    
    if instr_ptrn_from_datastack_file is not None: 
        ax.scatter(instr_peaks, instr_zeros, marker='x', color='r') 
    
    
    # plot element patterns 
        
    for i, ptrn in enumerate(ptrn_list):
        
        #plot_ptrn(ptrn['elem'], y, ax, color_select=color_select)
        
        plot_ptrn(ptrn['elem'], offset + i, ax, color_select=color_select)
                    
    ax.set_yticks(range(n_ticks))        
    ax.set_yticklabels(ytick_labels, fontsize=8) 
        
    ax.set_ylim([-1, n_ticks])
   
    
    #plt.tight_layout()
    
    return ax 
       

def plot_puzzle(datastack_file, n, elements=None, color_select=None, footspace=0.2): 
    '''Plot peak pattern puzzle for hotmax spectrum `n` with emission patterns. 
    
    If `elements` is specified, only plot patterns for those elements. 
    
    ''' 
    
    # generate patterns 
    if elements is None: 
        elements = eoi 
        
    elem_ptrns = get_patterns(elements=elements, color_select=color_select) 
    instrum_ptrn = get_instrument_pattern(datastack_file) 

    # prepare figure          
    n_ptrns = len(elements) + 1 
    n_grid_rows = n_ptrns + 12 
    
    fig_width = 8 
    fig_height = fig_width * n_grid_rows / 50 + 2
    
    fig = plt.figure(figsize=[fig_width, fig_height], constrained_layout=True)
    grid = plt.GridSpec(n_grid_rows, 1, hspace=16) 

    ax_ptrns = fig.add_subplot(grid[0:n_ptrns, 0])# , xticklabels=[])
    ax_spectr = fig.add_subplot(grid[n_ptrns+1:, 0], sharex=ax_ptrns) 
    
    # hack to suppress ticklabels in upper plot 
    # https://stackoverflow.com/questions/4209467 
    # matplotlib-share-x-axis-but-dont-show-x-axis-tick-labels-for-both-just-one
    plt.setp(ax_ptrns.get_xticklabels(), visible=False) 
    
    # make subplots 
    plot_patterns(elements=elements, color_select=color_select, instr_ptrn_from_datastack_file=datastack_file, 
                  ax=ax_ptrns) 
    
    #ax = plot_patterns(elem_ptrns, instrument_pattern=instrum_ptrn, ax=ax_ptrns) 
    
    hma = HotmaxAtlas(datastack_file) 
    hma.plot_spectrum(n, ax=ax_spectr, footspace=footspace, hotlines_ticklabels=False, headspace=1.3)

    ax_spectr.set_xlabel('Energy (keV)')
    
    # add tick labels explicitly to avoid issues with panning 
    _add_hotlines_ticklabels(datastack_file, ax_spectr)
    _add_hotlines_ticklabels(datastack_file, ax_ptrns, clip_vline=False)
    
    return fig, ax_ptrns, ax_spectr 
    


def get_instrument_pattern(datastack_file):
    '''Generate instrument peak pattern. 
    
    Pattern dictionary contains strongest Rhodium anode emission peaks,  
    their corresponding Compton shifted peaks, and a sensor peak 
    
    Returns: instrument_pattern_dict 
    
    '''
    
    ds = maxrf4u.DataStack(datastack_file) 
    
    x_keVs = ds.read('maxrf_energies')
    keV0 = ds.read('compton_peak_energy')[0]
    keV1 = maxrf4u.RHODIUM_Ka 
    theta = maxrf4u.detector_angle(keV0, keV1) 
    
    sensor_peak_idx = ds.read('hotmax_pixels')[0, 2]
    sensor_peak_keV = x_keVs[sensor_peak_idx]
    

    anode_emission = mos.XFluo('Rh', tube_keV=30, min_prom=0.1) # only largest K peaks  
    keV_in = anode_emission.peak_energies
    keV_out = maxrf4u.compton_shift(keV_in, theta)

    instrument_peaks = np.r_[sensor_peak_keV, keV_in, keV_out]
    instrument_peaks = np.sort(instrument_peaks) 
    
    instrument_pattern_dict = {'name': 'INSTRUMENT', 
                               'instrument_peaks': instrument_peaks}
    
    return instrument_pattern_dict
    

def add_hotlines(datastack_file, ax, vlines=True, clip_vline=True): 
    '''Utility function. Adds hotlines and tick labels to plot `ax`. 
    
    Instrument peaks are colored black.  
    
    '''
    # read hotlines data 
    ds = maxrf4u.DataStack(datastack_file) 
    x_keVs = ds.read('maxrf_energies')
    peak_idxs = ds.read('hotmax_pixels')[:, 2] 
    
    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 
