"""
@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.


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 interpolation at finer steps,
whereas the Kalha data contains only the original data points 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 requires python 3.6, numpy and the periodictable package (https://pypi.python.org/pypi/periodictable).

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

# read any periodictable's element interfaces as follows.
# eph and cs are numpy arrays of identical shape that hold the photon energies and cross sections.
eph, cs = pt.gold.photoionization.cross_section['4f']
eph, cs = pt.elements.symbol('Au').photoionization.cross_section['4f']
eph, cs = pt.elements.name('gold').photoionization.cross_section['4f']
eph, cs = pt.elements[79].photoionization.cross_section['4f']

# interpolate for specific photon energy
print(np.interp(photon_energy, eph, cs)
~~~~~~

the data is loaded from the cross-sections.dat file which is a python-pickled data file.
to switch between data sources, use one of the load functions defined here
and dump the data to the cross-sections.dat file.


@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
"""

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


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] = (eph, a[col].copy())
            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: (photon_energy, cross_section) tuple of 1-dimensional numpy arrays.
    """
    url = f"https://vuo.elettra.eu/services/elements/data/{symbol.lower()}{nl}.txt"
    try:
        data = urllib.request.urlopen(url)
    except urllib.error.HTTPError:
        eph = None
        cs = None
    else:
        a = np.genfromtxt(data)
        try:
            eph = a[:, 0]
            cs = a[:, 1]
        except IndexError:
            eph = None
            cs = None

    return eph, cs


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:
                eph, cs = load_elettra_file(element.symbol, nl)
                if eph is not None and cs is not None:
                    element_data[nl] = (eph, cs)
        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


class Photoionization(object):
    def __init__(self):
        self.cross_section = {}
        self.cross_section_units = "Mb"


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():
        try:
            el = table[int(el_key)]
        except ValueError:
            el = table.symbol(el_key)
        pi = Photoionization()
        pi.cross_section = el_data
        pi.cross_section_units = "Mb"
        el.photoionization = pi
