"""
@package pmsco.elements.spectrum
photoelectron spectrum simulator

this module calculates the basic structure of a photoelectron spectrum.
it calculates positions and approximate amplitude of elastic peaks
based on photon energy, binding energy, photoionization cross section, and stoichiometry.
escape depth, photon flux, analyser transmission are not accounted for.


usage
-----

this module requires python 3.6, numpy, matplotlib and
the periodictable package (https://pypi.python.org/pypi/periodictable).

~~~~~~{.py}
import numpy as np
import periodictable as pt
import pmsco.elements.spectrum as spec

# for working with the data
labels, energy, intensity = spec.build_spectrum(800., {"Ti": 1, "O": 2})

# for plotting
spec.plot_spectrum(800., {"Ti": 1, "O": 2})
~~~~~~



@author Matthias Muntwiler

@copyright (c) 2020 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

from matplotlib import pyplot as plt
import numpy as np
import periodictable as pt
from . import bindingenergy
from . import photoionization


def get_element(number_or_symbol):
    """
    return the given Element object of the periodic table.

    @param number_or_symbol: atomic number (int) or chemical symbol (str).
    @return: Element object.
    """
    try:
        el = pt.elements[number_or_symbol]
    except KeyError:
        el = pt.elements.symbol(number_or_symbol)
    return el


def get_binding_energy(photon_energy, element, nlj):
    """
    look up the binding energy of a core level and check whether it is smaller than the photon energy.

    @param photon_energy: photon energy in eV.
    @param element: Element object of the periodic table.
    @param nlj: (str) spectroscopic term, e.g. '4f7/2'.
    @return: (float) binding energy or numpy.nan.
    """
    try:
        eb = element.binding_energy[nlj]
    except KeyError:
        return np.nan
    if eb < photon_energy:
        return eb
    else:
        return np.nan


def get_cross_section(photon_energy, element, nlj):
    """
    look up the photoionization cross section.

    since the Yeh/Lindau tables do not resolve the spin-orbit splitting,
    this function applies the normal relative weights of a full sub-shell.

    the result is a linear interpolation between tabulated values.

    @param photon_energy: photon energy in eV.
    @param element: Element object of the periodic table.
    @param nlj: (str) spectroscopic term, e.g. '4f7/2'.
    @return: (float) cross section in Mb.
    """
    nl = nlj[0:2]
    try:
        pet, cst = element.photoionization.cross_section[nl]
    except KeyError:
        return np.nan

    # weights of spin-orbit peaks
    d_wso = {"p1/2": 1./3.,
             "p3/2": 2./3.,
             "d3/2": 2./5.,
             "d5/2": 3./5.,
             "f5/2": 3./7.,
             "f7/2": 4./7.}
    wso = d_wso.get(nlj[1:], 1.)
    cst = cst * wso

    # todo: consider spline
    return np.interp(photon_energy, pet, cst)


def build_spectrum(photon_energy, elements, binding_energy=False, work_function=4.5):
    """
    calculate the positions and amplitudes of core-level photoemission lines.

    the function looks up the binding energies and cross sections of all photoemission lines in the energy range
    given by the photon energy and returns an array of expected spectral lines.

    @param photon_energy: (numeric) photon energy in eV.
    @param elements: list or dictionary of elements.
        elements are identified by their atomic number (int) or chemical symbol (str).
        if a dictionary is given, the (float) values are stoichiometric weights of the elements.
    @param binding_energy: (bool) return binding energies (True) rather than kinetic energies (False, default).
    @param work_function: (float) work function of the instrument in eV.
    @return: tuple (labels, positions, intensities) of 1-dimensional numpy arrays representing the spectrum.
        labels are in the format {Symbol}{n}{l}{j}.
    """
    ekin = []
    ebind = []
    intens = []
    labels = []

    for element in elements:
        el = get_element(element)
        for n in range(1, 8):
            for l in "spdf":
                for j in ['', '1/2', '3/2', '5/2', '7/2']:
                    nlj = f"{n}{l}{j}"
                    eb = get_binding_energy(photon_energy, el, nlj)
                    cs = get_cross_section(photon_energy, el, nlj)
                    try:
                        cs = cs * elements[element]
                    except (KeyError, TypeError):
                        pass
                    if not np.isnan(eb) and not np.isnan(cs):
                        ekin.append(photon_energy - eb - work_function)
                        ebind.append(eb)
                        intens.append(cs)
                        labels.append(f"{el.symbol}{nlj}")

    ebind = np.array(ebind)
    ekin = np.array(ekin)
    intens = np.array(intens)
    labels = np.array(labels)

    if binding_energy:
        return labels, ebind, intens
    else:
        return labels, ekin, intens


def plot_spectrum(photon_energy, elements, binding_energy=False, work_function=4.5, show_labels=True):
    """
    plot a simple spectrum representation of a material.

    the function looks up the binding energies and cross sections of all photoemission lines in the energy range
    given by the photon energy and returns an array of expected spectral lines.

    the spectrum is plotted using matplotlib.pyplot.stem.

    @param photon_energy: (numeric) photon energy in eV.
    @param elements: list or dictionary of elements.
        elements are identified by their atomic number (int) or chemical symbol (str).
        if a dictionary is given, the (float) values are stoichiometric weights of the elements.
    @param binding_energy: (bool) return binding energies (True) rather than kinetic energies (False, default).
    @param work_function: (float) work function of the instrument in eV.
    @param show_labels: (bool) show peak labels (True, default) or not (False).
    @return: (figure, axes)
    """
    labels, energy, intensity = build_spectrum(photon_energy, elements, binding_energy=binding_energy,
                                               work_function=work_function)

    fig, ax = plt.subplots()
    ax.stem(energy, intensity, basefmt=' ', use_line_collection=True)
    if show_labels:
        for sxy in zip(labels, energy, intensity):
            ax.annotate(sxy[0], xy=(sxy[1], sxy[2]), textcoords='data')

    ax.grid()
    if binding_energy:
        ax.set_xlabel('binding energy')
    else:
        ax.set_xlabel('kinetic energy')
    ax.set_ylabel('intensity')
    ax.set_title(elements)
    return fig, ax