"""The forest and the trees"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/02_a-peak-pattern-atlas.ipynb.

# %% auto 0
__all__ = ['PeakPatternAtlas', 'colorize', 'PeriodicTable']

# %% ../notebooks/02_a-peak-pattern-atlas.ipynb 11
# Periodic Table 
import mendeleev as mv 
from mendeleev.fetch import fetch_table  

import moseley as mos 
from . import ElementXRF

# plotting 
import pandas as pd 
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches 
import matplotlib.colors as mcolors 

import holoviews as hv  


# %% ../notebooks/02_a-peak-pattern-atlas.ipynb 12
class PeakPatternAtlas():
    '''Create a PeakPatternAtlas instance. '''

    def __init__(self, EOI=None, excitation_energy_keV=25, x_keVs=None, verbose=True, order='alpha'):
        '''Compute an ordered Peak Pattern Atlas for `excitation_energy_keV=25`.'''

        self.table = mos.PeriodicTable(EOI=EOI)

        self.EOI = [element for element in self.table.EOI if not '#'in element]
        
        self.element_xrf_list = [] 
        self.ptrn_dict_list = []

        for i, element in enumerate(self.EOI):     
            if verbose: 
                print(f'Please wait while computing spectral pattern for element {i+1}/{len(self.EOI)}...', end='\r') 

            element_xrf = ElementXRF(element, excitation_energy_keV=excitation_energy_keV, x_keVs=x_keVs)
            ptrn_dict = element_xrf.ptrn_dict

            self.element_xrf_list.append(element_xrf)
            self.ptrn_dict_list.append(ptrn_dict)

        if order == 'alpha':
            alphas = [ptrn['alpha_keV'] for ptrn in self.ptrn_dict_list]
            idxs = np.argsort(alphas)
            
        elif order == 'Z': 
            atomic_numbers = [ptrn['atomic_number'] for ptrn in self.ptrn_dict_list] 
            idxs = np.argsort(atomic_numbers) 

        self.ptrn_dict_list = [self.ptrn_dict_list[idx] for idx in idxs]  
        self.element_xrf_list = [self.element_xrf_list[idx] for idx in idxs]  
        self.yticks = [(i, xrf.element) for i, xrf in enumerate(self.element_xrf_list)]
            
        if verbose: 
            print('Ready building Peak Pattern Atlas!                                   ') 

    def plot_atlas(self, moseley=False): 
        '''Build a holoviews peak pattern atlas. '''
    
        # standard element colors
        ptable = mos.PeriodicTable()
        color_dict = mos.colorize(ptable.ptable_regular)
        
        scatters = []
        lines = []
        yticks = []
        alphas = []
        
        Energy = hv.Dimension('Energy', unit='keV', range=(0, 20))
        
        for i, ptrn in enumerate(self.ptrn_dict_list): 
            # we should avoid recalculation here and use the ppa instance 
        
            element = ptrn['elem'] 
            Z = ptrn['atomic_number'] 
            color = color_dict[element]
            
            if moseley: 
                y = Z 
                color='grey'
            else: 
                y = i
                
            peaks_x, peaks_y = ptrn['peaks_xy'].T 
            x_min = float(peaks_x.min()) 
            x_max =float(peaks_x.max())
            ptrn_y = np.ones_like(peaks_y) * y
            
            alpha_keV = ptrn['alpha_keV'] 
                    
            scatters.append(hv.Scatter((peaks_x, ptrn_y), Energy, 'Element').opts(size=5, color=color))
            lines.append(hv.Curve([(x_min, y), (x_max, y)], Energy, 'Element').opts(color=color)) 
            alphas.append(hv.Scatter(([alpha_keV], [y]), Energy, 'Element').opts(size=7, marker='square', color=color)) 
        
            yticks.append((y, element))
        
        atlas = (hv.Overlay(lines) * hv.Overlay(scatters) * hv.Overlay(alphas))
        atlas = atlas.opts(yticks=yticks, frame_height=500, frame_width=500, xlim=(0, 22))
    
        return atlas
      


def _fetch_table(): 
    '''Utility function wrapper for mendeleev periodic table data.  
    
    Make dataframe with selected element attributes for regular part of the periodic table.
    The irregular Lanthanides and Actinides series are rare, so we do not plot them. 
    
    Returns: ptable_regular (dataframe) 
    '''
    
    #ptable_df = mv.get_table('elements') # mendeleev 0.5 
    ptable_df = fetch_table('elements')
    ptable = ptable_df[['atomic_number', 'symbol', 'name', 'group_id', 'period']]
    is_regular = ~np.isnan(ptable_df['group_id'].values)

    ptable_regular = ptable[is_regular]
    
    return ptable_regular 


def colorize(ptable_regular): 
    '''Generate regular element colors dictionary. 
    
    Returns: element_colors_dict '''
    
    n_regular = len(ptable_regular) # 90 rows 
    
    # see: https://matplotlib.org/stable/gallery/color/named_colors.html 
    css4_dict = mcolors.CSS4_COLORS 
        
    # hack to select brightest colors and exclude american double grays 
    css4_names = list(css4_dict.keys()) 
    css4_names = sorted(sorted(css4_names, key=lambda c: tuple(mcolors.rgb_to_hsv(mcolors.to_rgb(c))))) # brightest colors 
    css4_names = [name for name in css4_names if 'gray' not in name] # exclude doubles grays 

    # pick brightest 90 
    n_names = len(css4_names) 
    css4_names = css4_names[n_names - n_regular - 1:n_names - 1]

    element_colors_dict = dict(zip(ptable_regular['symbol'], css4_names))

    return element_colors_dict


def _draw_box(ax, element_attrs, edgecolor=None, facecolor=None):
    '''Draw a box for element *element_attrs*  in subplot *ax*. '''
    
    w = 0.9
    h = 0.9
    
    atomic_number, symbol, name, x, y = element_attrs 
    
    if edgecolor is None: 
        edgecolor = [0.8, 0.8, 1.0]
        
    if facecolor is None: 
        facecolor = [0.95, 0.95, 1.0]
    
    rect = mpatches.Rectangle([x - w/2, y - h/2], h, w, edgecolor=edgecolor, facecolor=facecolor)
    ax.add_patch(rect)
    
    ax.annotate(symbol, [x, y], xytext=[0, -4], textcoords='offset points',
                va='center', ha='center', fontsize=17)
    
    ax.annotate(atomic_number, [x, y], xytext=[0, 12], textcoords='offset points', 
                va='center', ha='center', fontsize=10)
    
    ax.annotate(name, [x, y], xytext=[0, -17], textcoords='offset points',
                va='center', ha='center', fontsize=5)


def _sort_ptrn_list(ptrn_dict_list, order='alpha'): 
    '''Utility function to sort ptrn_dict_list in 'alpha' or 'Z' order. 
    
    Returns: sorted_ptrn_list '''

    if order == 'alpha': 
        alpha_list = [ptrn['alpha_keV'] for ptrn in ptrn_dict_list]
        idxs = np.argsort(alpha_list)
    elif order == 'Z': 
        Z_list = [ptrn['atomic_number'] for ptrn in ptrn_dict_list] 
        idxs = np.argsort(Z_list)
        
    sorted_ptrn_list = [ptrn_dict_list[i] for i in idxs]  
        
    return sorted_ptrn_list


class PeriodicTable(): 
    '''Create a periodic table of elements. '''
    
    
    def __init__(self, EOI=None): 

        if EOI is None: 
            self.EOI = ['#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']    
        else: 
            self.EOI = EOI 
     
    
        # fetch regular part of periodic table 
        self.ptable_regular = _fetch_table() 
    
        # generate element colors dict 
        self.element_colors_dict = colorize(self.ptable_regular) 


    def plot_table(self, ax=None, figname=None): 
        '''Create periodic table plot with selected elements colorized.'''
    
        if ax is None: 
            fig, ax = plt.subplots(figsize=[14, 8])
        else: 
            fig = ax.get_figure()
            
        ax.invert_yaxis()
    
        for i, element_attrs in enumerate(self.ptable_regular.values): 
    
            atomic_number, symbol, name, x, y = element_attrs 
            facecolor = self.element_colors_dict[symbol] 
            
            if symbol in self.EOI: 
                edgecolor = 'r' 
            else: 
                edgecolor = None
            _draw_box(ax, element_attrs, edgecolor=edgecolor, facecolor=facecolor)
    
        ax.set_xlim(0, 19)
        ax.set_ylim(8, 0)
        ax.axis('off');
        
        if figname is not None: 
    
            fig.savefig(figname)
    
        return ax
