"""Predicting elemental XRF (X-Ray Fluorescence) spectra"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/01_theoretical-peak-patterns.ipynb.

# %% auto 0
__all__ = ['ElementXRF', 'gaussian_convolve', 'find_peaks', 'plot_pattern', 'get_attenuation']

# %% ../notebooks/01_theoretical-peak-patterns.ipynb 9
import numpy as np 
import xraydb 
import pandas as pd 
import scipy.signal as ssg 
import scipy.interpolate as sip
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import mendeleev

# %% ../notebooks/01_theoretical-peak-patterns.ipynb 10
class ElementXRF(): 
    '''Computes fluorescence emission line energies and intensities for `element`. 
    
    '''
    
    def __init__(self, element, excitation_energy_keV='rhodium_Ka', h_mm=0.1, x_keVs=None, std=0.01, min_prom=0.001): 
        '''Create a theoretical elemental xrf spectrum object. '''

        ### LINES ### 

        if excitation_energy_keV == 'rhodium_Ka': 
            excitation_energy_keV = 22.1  

        self.element = element 

        excitation_energy = 1000 * excitation_energy_keV # keV -> eV 

        lines = xraydb.xray_lines(element, excitation_energy=excitation_energy) 

        line_names = [] # Siegbahn notation
        line_labels = [] # IUPAC notation 
        line_energies = [] 
        line_probabilities = [] 

        line_fluo_yields = []
        line_jump_coeffs = [] 
        line_attenuations = [] 
        
        line_intensities = []
        
        for name, line in lines.items(): 

            # xraydb 'intensities' (a.k.a. transition probablities) sum up to unity within each level 
            emission_energy, probability, initial_level, final_level = line  
            emission_energy_keV = emission_energy / 1000  
            label = f'{element}_{initial_level}{final_level}' #  construct IUPAC notation 
            
            # get corresponding edge properties 
            edge = initial_level # IUPAC notation!  e.g. 'L1', not 'La'
            edge_energy, fluo_yield, jump_ratio = xraydb.xray_edge(element, edge) 
            jump_coeff = (jump_ratio - 1) / jump_ratio # see Volker Thomsen 

            attenuation = get_attenuation(element, emission_energy_keV, excitation_energy_keV, h_mm=h_mm)

            # multiplying edge jump coefficient, transition probability and fluorescence yield... 
            intensity = jump_coeff * probability * fluo_yield * attenuation

            line_names.append(name) 
            line_labels.append(label)
            line_energies.append(emission_energy_keV)
            line_probabilities.append(probability)  
            
            line_fluo_yields.append(fluo_yield)
            line_jump_coeffs.append(jump_coeff) 
            line_attenuations.append(attenuation) 
            
            line_intensities.append(intensity) 

        self.lines_table = pd.DataFrame({'name': line_names, 
                                         'label': line_labels, 
                                         'energy': line_energies, 
                                         'trans_prob': line_probabilities, 
                                         'fluo_yield': line_fluo_yields, 
                                         'jump_coeff': line_jump_coeffs, 
                                         'attenuation': line_attenuations, 
                                         'intensity': line_intensities}) 

    
        ### CONTINUOUS ### 
        
        self.spectrum_xy = gaussian_convolve(line_energies, line_intensities, x_keVs=x_keVs, std=std) 
        self.peaks_xy = find_peaks(*self.spectrum_xy, min_prom=min_prom) 

        # sort high to low (just in case) 
        sort_idxs = np.argsort(self.peaks_xy[:, 1])[::-1]
        self.peaks_xy = self.peaks_xy[sort_idxs]

        # find nearest line Siegbahn peak labels 
        peaks_x, peaks_y = self.peaks_xy.T
        
        self.peak_labels = []
        for peak_keV in peaks_x:     
            line_idx = int(np.argmin((self.lines_table['energy'].values   - peak_keV)**2))
            line_label = self.lines_table['name'].values[line_idx]
            self.peak_labels.append(line_label)

        # create element pattern dict 

        alpha_keV = float(self.peaks_xy[0][0])
        alpha_escape_keV = alpha_keV - 1.74  # Silicon detector escape energy shift 
        
        # element   
        elem = self.element
        name = mendeleev.element(elem).name
        atomic_number = mendeleev.element(elem).atomic_number
        
        self.ptrn_dict = {
            'elem': elem,  
            'atomic_number': atomic_number, 
            'alpha_keV': alpha_keV, 
            'name': name, 
            'peaks_xy': self.peaks_xy, 
            'peak_labels': self.peak_labels, 
            'alpha_escape_keV': alpha_escape_keV} 
        
        

    def plot_spectrum(self, ax=None, edgecolor=None, facecolor=None, vlines_colors=None, xlim=[0, 25], peak_labels=True): 
        
        if edgecolor is None: 
            edgecolor = cm.tab20(0)
        if facecolor is None: 
            facecolor = 'w' 
    
        if ax is None: 
            fig, ax = plt.subplots() 
    
        x, y = self.spectrum_xy 
        peaks_x, peaks_y = self.peaks_xy.T
    
        ax.plot(x, y, color=edgecolor, zorder=9)
        ax.scatter(peaks_x[0], peaks_y[0], marker='s', facecolor=facecolor, edgecolor=edgecolor, zorder=10)  
        ax.scatter(peaks_x[1:], peaks_y[1:], marker='o', facecolor=facecolor, edgecolor=edgecolor, zorder=10)

        # add peak labels
        if peak_labels: 
            for label, xy in zip(self.peak_labels, self.peaks_xy): 
                ax.annotate(label, xy, xytext=(0, 5), textcoords="offset points",
                            va="bottom", ha="center") 
    
        # need to interpolate spectrum to add vlines 
    
        spectrum = sip.interp1d(x, y)
        line_keVs = self.lines_table['energy'] 
    
        line_heights = spectrum(line_keVs)
    
        ax.vlines(line_keVs, np.zeros_like(line_keVs), line_heights, colors=vlines_colors, alpha=0.5, zorder=8) 

        ax.set_title(f'Theoretical X-Ray Fluorescence for {self.element}') 
        ax.set_xlabel('Energy (keV)')
        ax.set_xlim(xlim)
        ax.set_ylim(ymax=1.1 * y.max())
        
        return ax 



def gaussian_convolve(peak_energies, peak_intensities, x_keVs=None, std=0.01): 
    '''Convolves line spectrum defined by `peak_energies` and `peak_intensities` 
    
    with a Gaussian peak shape.'''
    
    if x_keVs is None: 
        x_keVs = np.linspace(0, 40, 10000)

    y_spectrum = np.zeros_like(x_keVs) 

    for peak_energy, peak_intensity in zip(peak_energies, peak_intensities): 

        y_spectrum += peak_intensity * np.exp(-(1 / std) * (x_keVs - peak_energy)**2)
        
    return x_keVs, y_spectrum  

def find_peaks(x, y, min_prom=0.001):  
    '''Find peaks in spectrum `x`, `y` with minimal prominence `min_prom=`
    
    Returns: peaks_xy''' 
    
    # find the peaks (channel indexes) 
    peaks, props = ssg.find_peaks(y/y.max(), prominence=[min_prom, None])

    peaks_x = x[peaks]
    peaks_y = y[peaks] 
    
    # sort according to peak_height 
    sort_idxs = np.argsort(peaks_y)[::-1] 

    peaks_x = peaks_x[sort_idxs]
    peaks_y = peaks_y[sort_idxs]
    
    peaks_xy = np.c_[peaks_x, peaks_y]

    return peaks_xy 

def plot_pattern(ptrn_dict, ax=None, offset=0, color=None): 
    '''Plot a single element peak pattern. 
    ''' 
        
    element = ptrn_dict['elem']
    
    x, _ = ptrn_dict['peaks_xy'].T
    y = np.zeros_like(x) + offset 
    alpha_escape_keV = ptrn_dict['alpha_escape_keV']
         
    left_x = min(x)
    right_x = max(x) 

    if color is None: 
        facecolor = 'w'
        edgecolor = 'grey'
        linecolor = 'grey'
        linestyle = ':'

    else: 
        facecolor = color 
        edgecolor = 'k'
        linecolor = color
        linestyle = '-'

    if ax is None: 
        fig, ax = plt.subplots() 

    ax.scatter(x, y, s=15, facecolor=facecolor, edgecolor=edgecolor)
    ax.scatter(x[0], y[0], marker='s', facecolor=facecolor, edgecolor=edgecolor)
    ax.plot([left_x, right_x], [y[0], y[0]], color=edgecolor, linestyle=linestyle, zorder=-10)
    ax.scatter(alpha_escape_keV, y[0], marker='|', color=edgecolor) 
    ax.annotate(element, [right_x, y[0]], xytext=[5, -1], fontsize=8, color='k', 
                textcoords='offset points', ha='left', va='center') 
   
    return ax 


def get_attenuation(element, emission_energy_keV, excitation_energy_keV='rhodium_Ka', h_mm=0.1): 
    '''Attenuation factors for emissions at `emission_keVs` for pure sample with thickness `h_mm`. 

    See Thomsen_2007 equation 11. 
    
    Returns: att_list (array of floats)'''

    if excitation_energy_keV == 'rhodium_Ka': 
        excitation_energy_keV = 22.1 
        
    excitation_energy = excitation_energy_keV * 1000 
    emission_energy = emission_energy_keV * 1000 

    if emission_energy < 150: # eV 
        attenuation = 0 
        
    else: 
        mu_excit = xraydb.mu_elam(element, excitation_energy)
        mu_emiss = xraydb.mu_elam(element, emission_energy)
    
        # mathematically the thickness factor is defined as: 
        rho = xraydb.atomic_density(element)
        thickness_factor = 1 - np.exp(-rho * h_mm * 0.1 * (mu_excit + mu_emiss))
        
        attenuation = (mu_excit / (mu_excit + mu_emiss)) * thickness_factor 

    return attenuation


