'''
Module ediff.pxrd
-----------------
Calculation of powder X-ray diffraction patterns.    
'''

import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import convolve as spConvolve
from pymatgen.core.structure import Structure as pmStructure
from pymatgen.analysis.diffraction.xrd import XRDCalculator as pmXRDCalculator

class Crystal:
    '''
    Define crystal structure.
    
    Parameters
    ----------
    structure : structure object
        * The *structure object* is usually obtained from a CIF-file
          (CIF = Crystallographic Information File).
        * The *CIF files* are available in crystallographic databases,
          such as http://www.crystallography.net
        * The exact type of the object is *pymatgen.core.structure.Structure*
              * This indicates that the structure is read from CIF by means of
                pymatgen package (which works behind the sceenes).
              * This also means that the structure can be created by any
                other way available in pymatgen https://pymatgen.org/
              * Nevertheless, for common usage it is enough to
                read the structure from CIF, ignoring technical details.
    temp_factors : float or dictionary, optional, the default is 0.8
        * Temperature factors characterize thermal movement of atoms.
        * If a float value is given,
          all elements have this temperature factor value.
        * If a dictionary is given,
          the elements have the values defined in the dictionary;
          a sample input dictionary: `temp_factors = {'Na':1.2, 'Cl':1.1}`.
    '''
    
    def __init__(self, structure, temp_factors=0.8):
        '''
        * Initialize Crystal object.
        * The parameters are described above in class definition.
        '''
        # (1) Define structure
        if type(structure) == pmStructure:
            # Structure was generated by PyMatGen methods
            self.structure = structure
        else:
            # Structure is not from PyMatGen => we suppose CIF input
            self.structure = self.read_structure_from_CIF(structure)
        # (2) Define temperature factors
        if type(temp_factors) == dict:
            # Temperature factors were given as dictionary
            self.temp_factors = temp_factors
        else:
            # Temp.factors are not dictionary => we suppose a number
            # (the number is a universal temperature factor for all atoms
            self.temp_factors = self.get_elements_with_temp_factors(
                self.structure, temp_factors)
    
    @staticmethod
    def read_structure_from_CIF(CIF):
        '''
        Read crystall structure from CIF file.

        Parameters
        ----------
        CIF : str or path object
            * The filename of CIF file.
            * CIF = Crystallographic Information File
              contains information about crystal structure.
            * CIF files are usually obtained from crystallographic databases,
              such as http://www.crystallography.net

        Returns
        -------
        structure : structure object
            * The exact type of the object is
              *pymatgen.core.structure.Structure*.
            * Nevertheless, in the structure objects are usually
              created from CIF files (pymatgen is hidden for a common user).
         '''
        structure = pmStructure.from_file(CIF)
        return(structure)

    def get_elements(self, structure):
        '''
        Get a list of all elements, which are contained in given structure.

        Parameters
        ----------
        structure : structure object
            Typically, the structure objects are
            created by function *read_structure_from_CIF* above.

        Returns
        -------
        list_of_elements : list
            List with symbols of all elements
            that are present in given *structure*.
            
        Technical note
        --------------
        * The elements are obtained from *structure* object.
        * The structure object = pymatgen.core.structure.Structure.
        * The list of element names si obtained by a set of tricks
          specific to pymatgen-structure object.
        * These tricks were revealed by inspecting the original
          pymatgen.core.structure.Structure in Spyder,
          by means of shortcuts Ctrl+I (help) and Ctrl+G (go to code).
        * Common users do not have to use this function.
        * The function is just used internally, when setting
          temperature factors of the elements contained in given *structure*.
        '''
        list_of_atoms = []
        for site in structure:
            for sp, occu in site.species.items():
                list_of_atoms.append(sp.symbol)
        list_of_elements = np.unique(list_of_atoms)
        return(list_of_elements)

    def get_elements_with_temp_factors(self, structure, B=0.8):
        '''
        Get a dictionary, which contains symbols and temperature factors
        of all elements, which are present in given structure.

        Parameters
        ----------
        structure : structure object
            Typically, the structure objects are
            created by function *read_structure_from_CIF* above.
        B : float, optional, the default is 0.8
            

        Returns
        -------
        elements_with_temp_factors : dict
            Dictionary with symbols and temperature factors
            of all elements, which are present in given *structure*.
            
        Notes
        -----
        * The dictionary is in the format required by pymatgen package.
        * The pymatgen package is used for PXRD calculations,
          but these calculations are hidden from common users.
        * Common users do not use this function.
        * The function is just used internally, when setting
          temperature factors of the elements contained in given *structure*.
         '''
        elements = self.get_elements(structure)
        elements_with_temp_factors = dict.fromkeys(elements, B)
        return(elements_with_temp_factors)

class Experiment:
    '''
    Define experimental parameters.
    
    Parameters
    ----------
    wavelength : float
        Wavelength of the X-rays.
        Typical values are 1.54 A (CuKa) or 0.71 A (MoKa).
    two_theta_range: list/tuple of two floats
        Minimal and maximal diffraction angle;
        both values are TwoTheta angle in [deg] (for given *wavelength*).
    '''
    def __init__(self, wavelength, two_theta_range):
        '''
        * Initialize Experimental object.
        * The parameters are described above in class definition.
        '''
        self.wavelength = wavelength
        self.two_theta_range = two_theta_range

class PlotParameters:
    '''
    Define local+global parameters for plotting.
    
    Parameters
    ----------
    title : str
        Title of the plot.
    x_axis : str, 'TwoTheta','S','q' or 'dhkl', optional, default is 'q'
        Quantity for X-axis.
    rcParams : dict; optional, the default is empty dictionary {}
        The dictionary should have the format of mathplotlib.pyplot.rcParams.
        The argmument is passed to matplotlib.pyplot.rcParams.update.
        The initialization procedure creates some default rcParams.
        This argument can override the pre-defined parameters,
        which means that the default is created anyway
        and then it can be adjusted here by rcParams argument.
    '''
    
    def __init__(self, title=None, x_axis='q', xlim=None, rcParams={}):
        '''
        * Initialize PlotParameters object.
        * The parameters are described above in class definition.
        '''
        
        # Initialize basic parameters
        self.title = title
        self.x_axis = x_axis
        self.xlim = xlim
        self.rcParams = rcParams
        # Set global plot settings using rcParams
        self.set_default_rcParams(rcParams)

    @staticmethod
    def set_default_rcParams(my_rcParams={}):
        '''
        A class method defining global plot parameters (plt.rcParams).
        
        Parameters
        ----------
        my_rcParams : dictionary
            containing selected plt.rcParams keys)
            DESCRIPTION. The default is {}.

        Returns
        -------
        None, BUT it redefines global variable plt.rcParams!

        Notes:
        ------
        * This is a @classmethod (because it is used within the whole class)
          but it could be a @staticmethod as well (because it does not use
          cls variable in fact).
        * The method is employed in two ways:
            - standard usage of MCREEP package: default rcParams are used
              (and possibly modified) in objects of PlotParameters class
            - special usage of MCREEP (more figures, multiplots): default
              rcParams are used when definining the axes of (multiple)figures
        '''
        # (1) Set default rcParams
        # (Hardcoded, suitable default for standard plots
        plt.rcParams.update({
            'figure.figsize'     : (8/2.54,6/2.54),
            'figure.dpi'         : 500,
            'font.size'          : 7,
            'lines.linewidth'    : 0.8,
            'axes.linewidth'     : 0.6,
            'xtick.major.width'  : 0.6,
            'ytick.major.width'  : 0.6,
            'grid.linewidth'     : 0.6,
            'grid.linestyle'     : ':'})
        # (2) Update default with argument rcParams, if it was given
        # (User-defined in the main program, if necessary
        # (Useful namely for multiplots
        plt.rcParams.update(my_rcParams)

class PeakProfiles:
    '''
    Define profile of diffraction peaks.

    * This class is employed only as a namespace.
    * It contains three functions/definitions of diffratction peak profiles.
    '''

    def gaussian(X,m,s):
        '''
        Gaussian function (~ profile for PXRD calculation).
    
        Parameters
        ----------
        X : numpy array
            X-variable of Gaussian function.
        m : float
            Mean value = the center of Gaussian function.
        s : float
            Standard deviation = the width of Gaussian function.
    
        Returns
        -------
        NumPy array
            The array with Y-values of the function; Y = Gaussian(X).
        '''
        return( 1/(s*np.sqrt(2*np.pi)) * np.exp(-(X-m)**2/(2*s**2)) )
    
    def lorentzian(X,m,s):
        '''
        Lorentzian function (~ profile for PXRD calculation).
    
        Parameters
        ----------
        X : numpy array
            X-variable of Lorentzian function.
        m : float
            Mean value = the center of Lorentzian function.
        s : float
            Standard deviation = the width of Lorentzian function.
    
        Returns
        -------
        NumPy array
            The array with Y-values of the functin; Y = Lorenzian(X).
        '''
        return( 1/np.pi * s/((X-m)**2 + s**2) )
    
    def pseudo_voigt(X,m,s,n=0.5):
        '''
        Pseudo-Voigt function (~ profile for PXRD calculation).
    
        Parameters
        ----------
        X : numpy array
            X-variable of pseudo-Voigt function.
        m : float
            Mean value = the center of pseudo-Voigt function.
        s : float
            Standard deviation = the width of pseudo-Voigt function.
        n : float, optional, the default is 0.5
            Ratio between Gaussian and Lorentzian components.
            In the very first approximation, we can set n = 0.5.
            For more details, see *Notes* subsection below.
            
        Returns
        -------
        NumPy array
            The array with Y-values of the functin; Y = Lorenzian(X).
    
        Notes
        -----
        * Pseudo-Voigt function
          = linear combination of Gaussian and Lorentzian function.
        * In the 1st approximation:
            * 50/50 combination => n = 0.5
            * sg = sigma_Gaussian
              defined so that sg and s yielded the same FWHM
        * more details in:
            * https://en.wikipedia.org/wiki/Voigt_profile
            * https://lmfit.github.io/lmfit-py/builtin_models.html
        '''
        # (1) Calculate sg
        # (sg = sigma_Gaussian - see docstring above for more details
        sg = s/np.sqrt(2*np.log(2))
        # (2) Calculate Pseudo-Voigt function
        # (Pseudo-Voigt = n * Gausian + (1-n) * Lorenzian)
        pseudo_voigt = \
            n * PeakProfiles.gaussian(X,m,sg) \
            + (1-n) * PeakProfiles.lorentzian(X,m,s)
        # (3) Return calculated function    
        return(pseudo_voigt)

class PXRDcalculation:
    '''
    Define calculation of PXRD = powder X-ray diffraction pattern.
    
    Parameters
    ----------
    crystal : ediff.pxrd.Crystal object
        This object is usually prepared in advance
        as an instance of ediff.pxrd.Crystal class.
    experiment : ediff.pxrd.Experiment object
        This object is usually prepared in advance
        as an instance of ediff.pxrd.Experiment class.        
    plot_parameters : ediff.pxrd.PlotParameters object
        This object is usually prepared in advance
        as an instance of ediff.pxrd.PlotParameters class.
    peak_profile_sigma : float, optional, the default is 0.03
        Width of the calculated diffraction peaks.
        The default = 0.03 is suitable for most of common calculations
        and corresponds to the default in older PowderCell program (n*FWHM=7).
        For overlapping peaks it may be slightly decreased,
        while for more realistic diffractograms it may be slightly increased. 
    peak_profile_type : None or ediff.pxrd.PeakProfiles object, optional
        Profile of the calculated diffraction peaks.
        The default is PeakProfiles.pseudo_voigt.
        This default is suitable for common calculations
        and does not have to be changed (in great majority of cases).
    '''
  
    def __init__(self, crystal, experiment, plot_parameters,      
                 peak_profile_sigma = 0.03,
                 peak_profile_type = PeakProfiles.pseudo_voigt):
        '''
        Initialize PXRDcalcualtion object.
        The parameters are described above in class definition.
        '''
        self.crystal = crystal
        self.experiment = experiment
        self.plot_parameters = plot_parameters
        self.peak_profile_sigma = peak_profile_sigma
        self.peak_profile_type = peak_profile_type 
        self.diffractions = self.calculate_diffractions()
        self.diffractogram = self.calculate_diffractogram()

    def calculate_diffractions(self):
        # (1) Calculate intensities
        calculation = pmXRDCalculator(
            wavelength = self.experiment.wavelength,
            debye_waller_factors=(self.crystal.temp_factors))
        diffractions = calculation.get_pattern(
            self.crystal.structure,
            two_theta_range = self.experiment.two_theta_range)
        # (2) Transform diffractions: PyMatGen structure => pd.DataFrame
        diffractions_df = self.diffractions_to_dframe(diffractions)
        # (3) Return calculated diffractions
        # (pandas.DataFrame with cols: TwoTheta, h,k,l, dhkl, S, q, Intensity
        return(diffractions_df)
        
    def calculate_diffractogram(self):
        # (1) Base diffraction profile
        # = diffractogram with zero intensities
        # = interval theta_min to theta_max with step=0.001, zero intensities
        th1,th2 = th1,th2 = self.experiment.two_theta_range
        numpoints = (th2-th1) * 100 + 1
        diffractogram_2theta = np.linspace(th1, th2, numpoints, endpoint=True)
        diffractogram_intensity = np.zeros(len(diffractogram_2theta))
        df = pd.DataFrame(
            np.transpose([diffractogram_2theta, diffractogram_intensity]),
            columns=['TwoTheta','Intensity'])
        # (2) Line diffraction pattern
        # = diffractogram with line intensities
        # (2a) read diffractions, rounded to 3 dec.places ~ step 0.001
        diffractions = np.round(
            np.array(self.diffractions[['TwoTheta','Ihkl']]),2)
        # (2b) insert line diffractions into zero profile from step (1)
        for two_theta, intensity in diffractions:
            # calculate index = position of the diffraction in profile array
            index = int(round((two_theta - th1) * 100))
            # insert the diffraction intensity at calculated position
            df.Intensity[index] = intensity
        # (3) Simulated diffraction profile
        # = diffractogram with intensity profiles
        # (3a) Is peak_profile_type defined?
        # (Argument {peak_profile_type} has default value - pseudo_voight
        # (Therefore, it should be defined unless user specifies None value
        if self.peak_profile_type == None:
            print('Error: peak_profile_type is not defined!')
            sys.exit()
        # (3b) The peak_peak_profile_type is defined => calculate difractogram
        # Precalculate peak profile
        X = np.linspace(-10,10,2001,endpoint=True)
        peak_profile = self.peak_profile_type(X,0,self.peak_profile_sigma)    
        peak_profile = peak_profile/np.max(peak_profile)
        # Convolve Intensities with PeakProfile
        df.Intensity = spConvolve(df.Intensity,peak_profile, mode='same')
        # Normalize the calculated intensities
        # (The normalization is not strictly necessary,
        # (as spConvolve defined above keeps the original normalization,
        # (but the normalization to 1 is convenient for further processing.
        df.Intensity = df.Intensity/np.max(df.Intensity)
        # (4) Final diffraction profile
        # = diffraction profile with additional columns (S,q)
        # = original cols (TwoTheta, Intensity) => (TwoTheta,S,q,Intensity)
        df = self.add_diffraction_vectors_to_diffractogram(df)
        # (5) Return diffractogram
        return(df)
        
    def print_diffractions(self):
        '''
        Print the calculated diffractions to stdout.
        '''
        table = PXRDcalculation.dframe_to_table(self.diffractions)
        print(table)
            
    def save_diffractions(self, output_file):
        '''
        Save the calculated diffractions to *output_file*.
        
        Parameters
        ----------
        output_file : str
            Name of the output file.

        Returns
        -------
        None; the output is the list of diffractions in the *output_file*.
        '''
        try:
            with open(output_file, 'w') as fh:
                table = PXRDcalculation.dframe_to_table(self.diffractions)
                print(table, file=fh)
        except:
            print(f'Error saving diffractions to {output_file}!')
    
    def plot_diffractions(self, outfile=None):
        '''
        Plot the calculated diffractions.
        
        Parameters
        ----------
        outfile : str, optional, the default is none
            Name of the output file.
            If not given, the plot is just shown, but not saved.

        Returns
        -------
        * None; the output is the plot on the screen (and outfile).
        * Note that the plot shows just diffraction intensities, not profiles.
        * Use *plot_diffractogram* for diffractogram with intensity profiles.
        '''
        # (1) Make local copies of pre-defined parameters (just convenience)
        df      = self.diffractions
        x_axis  = self.plot_parameters.x_axis
        # (2) Prepare the plot
        plt.vlines(df[x_axis],0,df.Ihkl, lw=2)
        # (3) Set plot details
        # (external function, common to diffractions and diffractogram plots
        self.set_plot_details(x_axis)
        # (4) Save/show the plot...
        # In Spyder: plot is always shown
        # In CLI: plot is saved (if outfile is defined) or it is just shown
        if outfile != None: 
            plt.savefig(outfile) 
        else:
            plt.show()
        
    def plot_diffractions_with_indexes(self):
        '''
        Plot indexed diffractions.
        
        Important - interactive plots in Spyder
        ---------------------------------------
        * This function can create interactive plot.
            * In CLI (command line) - the plot is interactive automatically.
            * In Spyder IDE - default is to create non-interactive plots.
        * If you want the interactive plot in Spyder IDE, type in Console ...
            * BEFORE running (to switch on interactivity): %matplotlib qt
            * AFTER running (back to non-interactive mode): %matplotlib inline
        
        Parameters
        ----------
        * None; this function is to be called exclusively 
          as method of PXRDcalculation object.
        * If PXRDcalculation is properly initialized,
          additional parameters are not needed.

        Returns
        -------
        * None; the output is the plot in the screen.
        * In typical case, the plot is interactive
          so that the indexed diffractions could be expected in detail.
        
        Technical notes
        ---------------
        * The code below uses (sligthly modified) PyMatGen functions.
        * Reason: PyMatGen plotting works well with diffraction indexes.
          Re-programing of indexed plots would be difficult and useless...
        '''
        # (0) Redefine plot parameters
        # (so that they were optimized for the interactive plot
        plt.rcParams.update({
            'figure.figsize'  : (12/2.54,4/2.54),
            'figure.dpi'      : 200,
            'font.size'       : 6})
        # (1) Calculate diffractions
        # (re-calculation in order to get the original PyMatGen xrd object
        # (the xrd object can be employed to plot diffractions with indexes
        xrd = pmXRDCalculator(
            wavelength = self.experiment.wavelength,
            debye_waller_factors=(self.crystal.temp_factors))
        # (2) Create plot with indexes
        # (trick: we create figure and axes (fig,ax)
        # (and instruct PyMatGen to plot in pre-defined axes
        fig,ax = plt.subplots()
        xrd.get_plot(
            structure = self.crystal.structure,
            two_theta_range=self.experiment.two_theta_range,
            annotate_peaks='compact', ax=ax, fontsize=6)
        # (3) Ex-post changing line properties in plt.plot
        # https://stackoverflow.com/q/41709257
        for line in ax.get_lines():
            line.set_color('orange')
        ax.set_xlabel('2theta [deg]')
        ax.set_ylabel('Intensity')
        ax.grid()
        # (4) Set higher number of ticks on X-axis
        # Trick #1: major ticks => higher MaxNLocator to get dense ticks
        # Trick #2: minor ticks => import special class for them
        ax.xaxis.set_major_locator(plt.MaxNLocator(15, steps=[1,2,5,10] ))
        from matplotlib.ticker import AutoMinorLocator
        ax.xaxis.set_minor_locator(AutoMinorLocator(5))
        # (5) Final adjustments
        fig.tight_layout()
        fig.show()
        
    def print_diffractogram(self):
        self.print_diffractions()
        print('-----')
        print('* Just diffraction intensities, not the whole pattern.')
        print('* Reason: the whole diffraction pattern is too long.')
        print('* Note: save_diffractogram save the pattern to TXT-file.')
    
    def save_diffractogram(self, outfile):
        my_title = 'Calculated PXRD diffractogram\n'
        np.savetxt(
            outfile,
            np.array(self.diffractogram),
            fmt = ['%6.2f','%8.4f','%8.4f','%8.4f'],
            header = my_title + 
                'Columns: TwoTheta[deg], S[1/A], q[1/A], Intensity')
    
    def plot_diffractogram(self, outfile, x_axis='q'):
        # (1) Make local copies of pre-defined parameters (just convenience)
        df     = self.diffractogram
        x_axis = self.plot_parameters.x_axis
        # (2) Prepare the plot
        plt.plot(df[x_axis],df.Intensity)
        # (3) Set plot details
        # (external function, common to diffractions and diffractogram plots
        self.set_plot_details(x_axis)
        # (4) Save/show the plot...
        # In Spyder: plot is always shown
        # In CLI: plot is saved (if outfile is defined) or it is just shown
        if outfile != None: 
            plt.savefig(outfile) 
        else:
            plt.show()
        
    @staticmethod    
    def diffractions_to_dframe(intensities):
        h = []; k = []; l = []
        for i in intensities.hkls:
            h.append(i[0]['hkl'][0])
            k.append(i[0]['hkl'][1])
            l.append(i[0]['hkl'][2])
        df = pd.DataFrame( 
            np.transpose(
                [intensities.x, h, k, l, intensities.d_hkls, intensities.y]),
            columns=['TwoTheta','h','k','l','dhkl','Ihkl'])
        df.insert(loc=5, column='S', value=1/df.dhkl)
        df.insert(loc=6, column='q', value=2*np.pi*df.S)
        return(df)
    
    @staticmethod
    def dframe_to_table(dframe):
        # POZOR: parametr formatters je zaludny!
        # * musi to byt sada funkci, takze...
        #   NEfunguji string-formaty '%.3f' - nejsou to funkce
        #   NEfunguji f-stringy - nelze prazdne - f'{8.3f}' ani f'{:8.3f}'
        # * muzou se pouzit funkce str.format nebo lambda funkce
        #   to uz je skoro jednodussi str.format
        #   trik #0: obecne formaty nahrazovaciho pole ve str.format
        #       {}, {nazev_pole}, {nazev_pole:specifikace_formatu} ...
        #       pricemz v nasem pripade mame {:specifikace_formatu}
        #       takze odtud se tam vezme ta zdanlive mysticka dvojtecka
        #   trik #1: float M.N-format je bez carky '{:8.3f}'.format
        #       ale: float  .N-format je s carkou  '{:,.3f}'.format
        #   trik #2: NElze naformatovat obecny typ jako integer
        #     takze: pro obecne cislo misto NEfunkcniho {:3d} nutno {:3.0f}
        table = dframe.to_string(
            formatters={
                'TwoTheta' : '{:8.3f}'.format,
                'h'        : '{:3.0f}'.format,
                'k'        : '{:3.0f}'.format,
                'l'        : '{:3.0f}'.format,
                'dhkl'     : '{:7.3f}'.format,
                'S'        : '{:7.3f}'.format,
                'q'        : '{:7.3f}'.format,
                'Ihkl'     : '{:9.3f}'.format
            }
        )
        return(table)

    def add_diffraction_vectors_to_diffractogram(self, df):
        # Prepare wavelenght (just shortcut for convenience)
        wavelength = self.experiment.wavelength
        # Insert a column with the calculated values of diffration vector S
        df.insert(loc = 1,
                  column = 'S',
                  value = 2*np.sin(df.TwoTheta/2*np.pi/180) / wavelength)
        # Insert a column with the calculated values of difraction vector q
        df.insert(loc = 2,
                  column = 'q',
                  value = 2*np.pi * df.S)
        # Return the modified/extended DataFrame
        return(df)

    def set_plot_details(self, x_axis):
        # (1) Make local copies of pre-defined parameters (just convenience)
        title   = self.plot_parameters.title
        x_axis  = self.plot_parameters.x_axis
        my_xlim = self.plot_parameters.xlim
        # (2) Set plot title (only if it is defined)
        if title != None: plt.title(title)
        # (3) Set xlimits = xlim (only if they are defined)
        if my_xlim != None: plt.xlim(my_xlim)
        # (4) Determine label of x-axis = xlabel; it depends on x_axis argument
        if (x_axis == 'q') or (x_axis == 'S'):
            xlabel_str = f'${x_axis}$' + r' [1/$\mathrm{\AA}$]'
        elif x_axis == 'TwoTheta':
            xlabel_str = r'2$\theta$ [deg]'
        elif x_axis == 'dhkl':
            xlabel_str = r'$d(hkl)$' + r' [$\mathrm{\AA}$]'
        else:
            print('X-axis - unknown quantity and units...')
        # (4) Set the remaining plot details
        plt.xlabel(xlabel_str)
        plt.ylabel('Intensity')
        plt.grid()
        plt.tight_layout()
