"""
@package pmsco.elements.photoionization
Photoionization cross-sections of the elements

Extends the element table of the `periodictable` package
(https://periodictable.readthedocs.io/en/latest/index.html)
by a table of photoionization cross-sections and asymmetry parameters.


The data is available from (https://vuo.elettra.eu/services/elements/)
or (https://figshare.com/articles/dataset/Digitisation_of_Yeh_and_Lindau_Photoionisation_Cross_Section_Tabulated_Data/12389750).
Both sources are based on the original atomic data tables by Yeh and Lindau (1985).
The Elettra data includes the cross section and asymmetry parameter and is interpolated at finer steps,
whereas the Kalha data contains only the cross sections at the photon energies calculated by Yeh and Lindau
plus an additional point at 8 keV.
The tables go up to 1500 eV photon energy and do not resolve spin-orbit splitting.


Usage
-----

This module adds the photoionization attribute to the elements database of the periodictable package (https://pypi.python.org/pypi/periodictable).
Python >= 3.6, numpy >= 1.15 and the periodictable package are required.

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

# get a SubShellPhotoIonization object from any of periodictable's element interface:
sspi = pt.gold.photoionization['4f']
sspi = pt.elements.symbol('Au').photoionization['4f']
sspi = pt.elements.name('gold').photoionization['4f']
sspi = pt.elements[79].photoionization['4f']

# get the cross section, asymmetry parameter or differential cross section at 800 eV photon energy:
sspi.cross_section(800)
sspi.asymmetry_parameter(800)
sspi.diff_cross_section(800, gamma=30)

# with the j quantum number, the cross-section is weighted based on a full sub-shell:
sspi = pt.gold.photoionization['4f7/2']
print(sspi.weight)
print(pt.gold.photoionization['4f7/2'].cross_section(800) / pt.gold.photoionization['4f'].cross_section(800))

# the original data is contained in the data array (which is a numpy.recarray):
sspi.data.eph, sspi.data.cs, sspi.data.ap
~~~~~~

The data is loaded on demand from the cross-sections.dat file when the photoionization record is first accessed.
Normally, the user will not need to call any functions in this module directly.

The load_elettra_data()/load_kalha_data() and save_pickled_data() functions are provided
to import data from one of the sources referenced above and
to create the cross-sections.dat file.


@author Matthias Muntwiler

@copyright (c) 2020-23 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
"""

import copy
import numpy as np
from pathlib import Path
import periodictable as pt
import pickle
import urllib.request
import urllib.error
import periodictable.core


class PhotoIonization(dict):
    """
    photo-ionization parameters of an element

    this class provides the photo-ionization cross-section and asymmetry parameter of the sub-shells of an element.
    it is, essentially, a dictionary, mapping 'nl' and 'nlj' terms to the corresponding SubShellPhotoIonization object.

    examples of 'nl' and 'nlj' terms: '4f' and '4f7/2'

    @note the dictionary actually contains raw data for 'nl' terms only.
        for 'nlj' terms, the corresponding 'nl' object is copied,
        and a weight according to the spin-orbit multiplicity is set.

    @note 'nlj' terms are not considered by any methods or properties
        except the bracket notation or __getitem__ method!
        in particular, iteration or the keys() method will yield 'nl' terms only.
    """

    def __init__(self, *args, **kwargs):
        """
        dictionary constructor

        the class accepts the same arguments as the Python built-in dict constructor.
        keys are 'nl' terms, e.g. '4f', and values must be SubShellPhotoIonization() objects.

        @param args:
        @param kwargs:
        """
        super().__init__(*args, **kwargs)
        self.cross_section_units = "Mb"

    def __getitem__(self, k):
        """
        get sub-shell photo-ionization data by 'nl' or 'nlj' term.

        @param k: dictionary key.
            if this is an 'nl' term, the original object is returned.
            if this is an 'nlj' term, a proxy of the corresponding 'nl' object
            with shared data but weight based on j-branching is returned.

        @return: SubShellPhotoIonization() object

        @note whether the original or a proxy object is returned,
            its data attribute always refers to the original data.
            any modification will affect the original data (process memory).
        """
        spi = super().__getitem__(k[0:2])
        if len(k) > 2:
            spi = copy.copy(spi)
            spi.set_spin_orbit(k[1:5])
        return spi


class SubShellPhotoIonization(object):
    """
    Sub-shell photo-ionization parameters versus photon energy.

    this class provides the photo-ionization cross-section and asymmetry parameter of one sub-shell.
    it contains a three-column record array of photon energy, cross section and asymmetry parameter in self.data.
    accessory functions provide high-level access to specific views and interpolated data.

    a weighting factor self.weight is multiplied to the method results.
    it is normally used to weight the spin-orbit peaks by calling set_spin_orbit().
    """
    SPIN_ORBIT_WEIGHTS = {"p1/2": 1. / 3.,
                          "p3/2": 2. / 3.,
                          "d3/2": 2. / 5.,
                          "d5/2": 3. / 5.,
                          "f5/2": 3. / 7.,
                          "f7/2": 4. / 7.}

    def __init__(self, photon_energy, cross_section, asymmetry_parameter):
        """
        initialize a new object instance.

        all arrays must have the same length.

        @param photon_energy: (array-like) photon energies
        @param cross_section: (array-like) cross-section values
        @param asymmetry_parameter: (array-like) asymmetry parameter values
        """
        super().__init__()
        self.data = np.rec.fromarrays([photon_energy, cross_section, asymmetry_parameter], names='eph, cs, ap')
        self.weight = 1.

    def cross_section(self, photon_energy):
        """
        interpolated sub-shell cross-section at a specific energy.

        the weighting factor self.weight (e.g. spin-orbit) is included in the result.

        @param photon_energy: photon energy in eV.
            can be scalar or numpy array.
        @return: cross-section in Mb.
            numpy.nan where photon_energy is off range.
        """
        cs = np.interp(photon_energy, self.data.eph, self.data.cs, left=np.nan, right=np.nan) * self.weight
        return cs

    def asymmetry_parameter(self, photon_energy):
        """
        interpolated asymmetry parameter at a specific energy.

        @param photon_energy: photon energy in eV.
            can be scalar or numpy array.
        @return: asymmetry parameter (0..2).
            numpy.nan where photon_energy is off range.
        """
        ap = np.interp(photon_energy, self.data.eph, self.data.ap, left=np.nan, right=np.nan)
        return ap

    def diff_cross_section(self, photon_energy, gamma):
        """
        differential cross-section for linear polarization.

        the weighting factor self.weight (e.g. spin-orbit) is included in the result.

        @param photon_energy: photon energy in eV.
        @param gamma: angle between polarization vector and electron propagation direction in degrees.
        @return: differential cross-section in Mb.
        """
        p2 = (3 * np.cos(gamma) ** 2 - 1) / 2
        cs = self.cross_section(photon_energy)
        ap = self.asymmetry_parameter(photon_energy)
        dcs = cs / 4 / np.pi * (1 + ap * p2)
        return dcs

    def photon_energy_array(self):
        """
        photon energy array.

        the weighting factor self.weight (e.g. spin-orbit) is included in the result.

        @return:
        """
        return self.data.eph

    def cross_section_array(self):
        """
        sub-shell cross-section versus photon energy.

        the weighting factor self.weight (e.g. spin-orbit) is included in the result.

        @return: numpy.ndarray
        """
        return self.data.cs * self.weight

    def asymmetry_parameter_array(self):
        """
        sub-shell asymmetry parameter versus photon energy.

        the weighting factor self.weight (e.g. spin-orbit) is included in the result.

        @return: numpy.ndarray
        """
        return self.data.ap

    def diff_cross_section_array(self, gamma):
        """
        differential cross-section for linear polarization (full array).

        @param gamma: angle between polarization vector and electron propagation direction in degrees.
        @return: (np.ndarray) differential cross-section in Mb.
        """
        p2 = (3 * np.cos(gamma) ** 2 - 1) / 2
        dcs = self.data.cs / 4 / np.pi * (1 + self.data.ap * p2) * self.weight
        return dcs

    def set_spin_orbit(self, lj):
        """
        set the weight according to the spin-orbit quantum number (based on full sub-shell).

        the weight is stored in the self.weight attribute.
        it is applied to the results of the cross-section methods, but not to the raw data in self.data!

        @param lj: (str) 4-character lj term notation, e.g. 'f7/2'
        @return: None
        """
        self.weight = self.SPIN_ORBIT_WEIGHTS.get(lj, 1.)


def load_kalha_data():
    """
    load all cross-sections from csv-files by Kalha et al.

    the files must be placed in the 'kalha' directory next to this file.

    @return: cross-section data in a nested dictionary, cf. load_pickled_data().
    """
    data = {}
    p = Path(Path(__file__).parent, "kalha")
    for entry in p.glob('*_*.csv'):
        if entry.is_file():
            try:
                element = int(entry.stem.split('_')[0])
            except ValueError:
                pass
            else:
                data[element] = load_kalha_file(entry)
    return data


def load_kalha_file(path):
    """
    load the cross-sections of an element from a csv-file by Kalha et al.

    @param path: file path
    @return: (dict) dictionary of 'nl' terms.
        the data items are tuples (photon_energy, cross_sections) of 1-dimensional numpy arrays.
    """
    a = np.genfromtxt(path, delimiter=',', names=True)
    b = ~np.isnan(a['Photon_Energy__eV'])
    a = a[b]
    eph = a['Photon_Energy__eV'].copy()
    data = {}
    for n in range(1, 8):
        for l in 'spdf':
            col = f"{n}{l}"
            try:
                data[col] = SubShellPhotoIonization(eph, a[col].copy(), np.zeros_like(eph))
            except ValueError:
                pass
    return data


def load_kalha_configuration(path):
    """
    load the electron configuration from a csv-file by Kalha et al.

    @param path: file path
    @return: (dict) dictionary of 'nl' terms mapping to number of electrons in the sub-shell.
    """
    p = Path(path)
    subshells = []
    electrons = []
    config = {}
    with p.open() as f:
        for l in f.readlines():
            s = l.split(',')
            k_eph = "Photon Energy"
            k_el = "#electrons"
            if s[0][0:len(k_eph)] == k_eph:
                subshells = s[1:]
            elif s[0][0:len(k_el)] == k_el:
                electrons = s[1:]

    for i, sh in enumerate(subshells):
        if sh:
            config[sh] = electrons[i]

    return config


def load_elettra_file(symbol, nl):
    """
    download the cross sections of one level from the Elettra webelements web site.

    @param symbol: (str) element symbol
    @param nl: (str) nl term, e.g. '2p' (no spin-orbit)
    @return: PhotoIonizationData(photon_energy, cross_section, asymmetry_parameter)
        named tuple of 1-dimensional numpy arrays.
    """
    spi = None

    url = f"https://vuo.elettra.eu/services/elements/data/{symbol.lower()}{nl}.txt"
    try:
        data = urllib.request.urlopen(url)
    except urllib.error.HTTPError:
        pass
    else:
        a = np.genfromtxt(data)
        try:
            spi = SubShellPhotoIonization(a[:, 0], a[:, 1], a[:, 4])
        except IndexError:
            pass

    return spi


def load_elettra_data():
    """
    download the cross sections from the Elettra webelements web site.

    @return: cross-section data in a nested dictionary, cf. load_pickled_data().
    """
    data = {}
    for element in pt.elements:
        element_data = {}
        for nlj in element.binding_energy:
            nl = nlj[0:2]
            eb = element.binding_energy[nlj]
            if nl not in element_data and eb <= 2000:
                spi = load_elettra_file(element.symbol, nl)
                if spi is not None:
                    element_data[nl] = spi
        if len(element_data):
            data[element.symbol] = element_data

    return data


def save_pickled_data(path, data):
    """
    save a cross section data dictionary to a python-pickled file.

    @param path: file path
    @param data: cross-section data in a nested dictionary, cf. load_pickled_data().
    @return: None
    """
    with open(path, "wb") as f:
        pickle.dump(data, f)


def load_pickled_data(path):
    """
    load the cross section data from a python-pickled file.

    the file can be generated by the save_pickled_data() function.

    @param path: file path
    @return: cross-section data in a nested dictionary.
        the first-level keys are element symbols.
        the second-level keys are 'nl' terms (e.g. '2p').
        note that the Yeh and Lindau tables do not resolve spin-orbit splitting.
        the data items are (photon_energy, cross_sections) tuples
        of 1-dimensional numpy arrays holding the data table.
        cross section values are given in Mb.
    """
    with open(path, "rb") as f:
        data = pickle.load(f)
    return data


def init(table, reload=False):
    """
    loads cross-section data into the periodic table.

    this function is called by the periodictable to load the data on demand.

    @param table:
    @param reload:
    @return:
    """
    if 'photoionization' in table.properties and not reload:
        return
    table.properties.append('photoionization')

    # default value
    pt.core.Element.photoionization = PhotoIonization()

    p = Path(Path(__file__).parent, "cross-sections.dat")
    data = load_pickled_data(p)
    for el_key, el_data in data.items():
        # el_data is dict('nl': PhotoIonizationData)
        try:
            el = table[int(el_key)]
        except ValueError:
            el = table.symbol(el_key)
        el.photoionization = PhotoIonization(el_data)


def _load_photoionization():
    """
    delayed loading of the binding energy table.
    """

    init(periodictable.core.default_table())


periodictable.core.delayed_load(['photoionization'], _load_photoionization)
