"""
@package pmsco.calculators.phagen.translator
Natoli/Sebilleau PHAGEN interface

this module provides conversion between input/output files of PHAGEN and EDAC.

@author Matthias Muntwiler

@copyright (c) 2015-19 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 logging
import numpy as np

from pmsco.cluster import Cluster

logger = logging.getLogger(__name__)

## rydberg energy in electron volts
ERYDBERG = 13.6056923


def state_to_edge(state):
    """
    translate spectroscopic notation to edge notation.

    @param state: spectroscopic notation: "1s", "2s", "2p1/2", etc.
    @return: edge notation: "k", "l1", "l2", etc.
        note: if the j-value is not given, the lower j edge is returned.
    """
    jshells = ['s', 'p1/2', 'p3/2', 'd3/2', 'd5/2', 'f5/2', 'f7/2']
    lshells = [s[0] for s in jshells]
    shell = int(state[0])
    try:
        subshell = jshells.index(state[1:]) + 1
    except ValueError:
        subshell = lshells.index(state[1]) + 1
    except IndexError:
        subshell = 1
    edge = "klmnop"[shell-1]
    if shell > 1:
        edge += str(subshell)
    return edge


class TranslationParams(object):
    """
    project parameters needed for translation.

    energy unit is eV.
    """
    def __init__(self):
        self.initial_state = "1s"
        self.binding_energy = 0.
        self.cluster = None
        self.kinetic_energies = np.empty(0, dtype=float)

    @property
    def l_init(self):
        return "spdf".index(self.initial_state[1])

    @property
    def edge(self):
        return state_to_edge(self.initial_state)

    def set_params(self, params):
        """
        set the translation parameters.

        @param params: a pmsco.project.CalculatorParams object or
                       a dictionary containing some or all public fields of this class.
        @return: None
        """
        try:
            self.initial_state = params.initial_state
            self.binding_energy = params.binding_energy
        except AttributeError:
            for key in params:
                self.__setattr__(key, params[key])

    def set_scan(self, scan):
        """
        set the scan parameters.

        @param scan: a pmsco.project.Scan object
        @return: None
        """
        try:
            energies = scan.energies
        except AttributeError:
            try:
                energies = scan['e']
            except KeyError:
                energies = scan
        if not isinstance(energies, np.ndarray):
            energies = np.array(energies)
        self.kinetic_energies = np.resize(self.kinetic_energies, energies.shape)
        self.kinetic_energies = energies

    def set_cluster(self, cluster):
        """
        set the initial cluster.

        @param cluster: a pmsco.cluster.Cluster object
        @return: None
        """
        self.cluster = cluster


class Translator(object):
    """
    data conversion to/from phagen input/output files.

    usage:
    1. set the translation parameters self.params.
    2. call write_input_file to create the phagen input files.
    3. call phagen on the input file.
    4. call parse_phagen_phase.
    5. call parse_radial_file.
    6. call write_edac_scattering to produce the EDAC scattering matrix files.
    7. call write_edac_emission to produce the EDAC emission matrix file.
    """

    ## @var params
    #
    # project parameters needed for translation.
    #
    # fill the attributes of this object before using any translator methods.

    ## @var scattering
    #
    # t-matrix storage
    #
    # the t-matrix is stored in a flat, one-dimensional numpy structured array consisting of the following fields:
    # @arg e (float) energy (eV)
    # @arg a (int) atom index (1-based)
    # @arg l (int) angular momentum quantum number l
    # @arg t (complex) scattering matrix element, t = exp(-i * delta) * sin delta
    #
    # @note PHAGEN uses the convention t = exp(-i * delta) * sin delta,
    # whereas EDAC uses t = exp(i * delta) * sin delta (complex conjugate).
    # this object stores the t-matrix according to the PHAGEN convention.
    # the conversion to the EDAC convention occurs in write_edac_scattering_file().

    ## @var emission
    #
    # radial matrix element storage
    #
    # the radial matrix elemnts are stored in a flat, one-dimensional numpy structured array
    # consisting of the following fields:
    # @arg e (float) energy (eV)
    # @arg dw (complex) matrix element for the transition to l-1
    # @arg up (complex) matrix element for the transition to l+1

    ## @var cluster
    #
    # cluster object for PHAGEN
    #
    # this object is created by translate_cluster().

    def __init__(self):
        """
        initialize the object instance.
        """
        self.params = TranslationParams()
        dt = [('e', 'f4'), ('a', 'i4'), ('l', 'i4'), ('t', 'c16')]
        self.scattering = np.empty(0, dtype=dt)
        dt = [('e', 'f4'), ('dw', 'c16'), ('up', 'c16')]
        self.emission = np.empty(0, dtype=dt)
        self.cluster = None

    def translate_cluster(self):
        """
        translate the cluster into a form suitable for PHAGEN.

        specifically, move the (first and hopefully only) emitter to the first atom position.

        the method copies the cluster from self.params into a new object
        and stores it under self.cluster.

        @return: None
        """
        self.cluster = Cluster()
        self.cluster.copy_from(self.params.cluster)
        ems = self.cluster.get_emitters(['i'])
        self.cluster.move_to_first(idx=ems[0][0]-1)

    def write_cluster(self, f):
        """
        write the cluster section of the PHAGEN input file.

        @param f: file or output stream (an object with a write method)

        @return: None
        """
        for atom in self.cluster.data:
            d = {k: atom[k] for k in atom.dtype.names}
            f.write("{s} {t} {x} {y} {z}\n".format(**d))
        f.write("-1 -1 0. 0. 0.\n")

    def write_ionicity(self, f):
        """
        write the ionicity section of the PHAGEN input file.

        ionicity is read from the 'q' column of the cluster.
        all atoms of a chemical element must have the same charge state
        because ionicity has to be specified per element.
        this function writes the average of all charge states of an element.

        @param f: file or output stream (an object with a write method)

        @return: None
        """
        data = self.cluster.data
        elements = np.unique(data['t'])
        for element in elements:
            idx = np.where(data['t'] == element)
            charge = np.mean(data['q'][idx])
            f.write("{t} {q}\n".format(t=element, q=charge))

        f.write("-1\n")

    def write_input(self, f):
        """
        write the PHAGEN input file.

        @param f: file path or output stream (an object with a write method).

        @return: None
        """
        phagen_params = {}

        self.translate_cluster()
        phagen_params['absorber'] = 1
        phagen_params['emin'] = self.params.kinetic_energies.min() / ERYDBERG
        phagen_params['emax'] = self.params.kinetic_energies.max() / ERYDBERG
        if self.params.kinetic_energies.shape[0] > 1:
            phagen_params['delta'] = (phagen_params['emax'] - phagen_params['emin']) / \
                                     (self.params.kinetic_energies.shape[0] - 1)
        else:
            phagen_params['delta'] = 0.1
        phagen_params['edge'] = state_to_edge(self.params.initial_state)
        phagen_params['edge1'] = 'm4'  # auger not supported
        phagen_params['edge2'] = 'm4'  # auger not supported
        phagen_params['cip'] = self.params.binding_energy / ERYDBERG
        if phagen_params['cip'] < 0.001:
            raise ValueError("binding energy parameter is zero.")

        if np.sum(np.abs(self.cluster.data['q'])) > 0.:
            phagen_params['ionzst'] = 'ionic'
        else:
            phagen_params['ionzst'] = 'neutral'

        if hasattr(f, "write") and callable(f.write):
            f.write("&job\n")
            f.write("calctype='xpd',\n")
            f.write("coor='angs',\n")
            f.write("cip={cip},\n".format(**phagen_params))
            f.write("absorber={absorber},\n".format(**phagen_params))
            f.write("edge='{edge}',\n".format(**phagen_params))
            f.write("edge1='{edge1}',\n".format(**phagen_params))
            f.write("edge2='{edge1}',\n".format(**phagen_params))
            f.write("gamma=0.03,\n")
            f.write("lmax_mode=2,\n")
            f.write("lmaxt=50,\n")
            f.write("emin={emin},\n".format(**phagen_params))
            f.write("emax={emax},\n".format(**phagen_params))
            f.write("delta={delta},\n".format(**phagen_params))
            f.write("potgen='in',\n")
            f.write("potype='hedin',\n")
            f.write("norman='stdcrm',\n")
            f.write("ovlpfac=0.0,\n")
            f.write("ionzst='{ionzst}',\n".format(**phagen_params))
            f.write("charelx='ex',\n")
            f.write("l2h=4\n")
            f.write("&end\n")
            f.write("comment 1\n")
            f.write("comment 2\n")
            f.write("\n")

            self.write_cluster(f)
            self.write_ionicity(f)
        else:
            with open(f, "wt", encoding="latin1") as fi:
                self.write_input(fi)

    def parse_phagen_phase(self, f):
        """
        parse the phase output file from PHAGEN.

        the phase file is written to div/phases.dat.
        it contains the following columns:

        @arg e energy (Ry)
        @arg x1 unknown 1
        @arg x2 unknown 2
        @arg na atom index (1-based)
        @arg nl angular momentum quantum number l
        @arg tr real part of the scattering matrix element
        @arg ti imaginary part of the scattering matrix element
        @arg ph phase shift

        the data is translated into the self.scattering array.

        @arg e energy (eV)
        @arg a atom index (1-based)
        @arg l angular momentum quantum number l
        @arg t complex scattering matrix element

        @note PHAGEN uses the convention t = exp(-i * delta) * sin delta,
        whereas EDAC uses t = exp(i * delta) * sin delta (complex conjugate).
        this class stores the t-matrix according to the PHAGEN convention.
        the conversion to the EDAC convention occurs in write_edac_scattering_file().

        @param f: file or path (any file-like or path-like object that can be passed to numpy.genfromtxt).

        @return: None
        """
        dt = [('e', 'f4'), ('x1', 'f4'), ('x2', 'f4'), ('na', 'i4'), ('nl', 'i4'),
              ('tr', 'f8'), ('ti', 'f8'), ('ph', 'f4')]
        data = np.atleast_1d(np.genfromtxt(f, dtype=dt))

        self.scattering = np.resize(self.scattering, data.shape)
        scat = self.scattering
        scat['e'] = data['e'] * ERYDBERG
        scat['a'] = data['na']
        scat['l'] = data['nl']
        scat['t'] = data['tr'] + 1j * data['ti']

    def write_edac_scattering(self, filename_format, phases=False):
        """
        write scatterer files for EDAC.

        produces one file for each atom class in self.scattering.

        @param filename_format: file name including a placeholder {} for the atom class.

        @param phases: write phase files instead of t-matrix files.

        @return: dictionary that maps atom classes to file names
        """
        if phases:
            write = self.write_edac_phase_file
        else:
            write = self.write_edac_scattering_file
        scat = self.scattering
        atoms = np.unique(scat['a'])
        files = {}
        for atom in atoms:
            f = filename_format.format(atom)
            sel = scat['a'] == atom
            idx = np.where(sel)
            atom_scat = scat[idx]
            write(f, atom_scat)
            files[atom] = f

        return files

    def write_edac_scattering_file(self, f, scat):
        """
        write a scatterer file for EDAC.

        @param f: file path or output stream (an object with a write method).

        @param scat: a slice of the self.scattering array belonging to the same atom class.

        @return: None
        """
        if hasattr(f, "write") and callable(f.write):
            energies = np.unique(scat['e'])
            ne = energies.shape[0]
            lmax = scat['l'].max()
            if ne == 1:
                f.write("1 {lmax} regular tl\n".format(lmax=lmax))
            else:
                f.write("{nk} E(eV) {lmax} regular tl\n".format(nk=ne, lmax=lmax))
            for energy in energies:
                sel = scat['e'] == energy
                idx = np.where(sel)
                energy_scat = scat[idx]
                if ne > 1:
                    f.write("{0:.3f} ".format(energy))
                for item in energy_scat:
                    f.write(" {0:.6f} {1:.6f}".format(item['t'].real, -item['t'].imag))
                for i in range(len(energy_scat), lmax + 1):
                    f.write(" 0 0")
                f.write("\n")
        else:
            with open(f, "wt", encoding="latin1") as fi:
                self.write_edac_scattering_file(fi, scat)

    def write_edac_phase_file(self, f, scat):
        """
        write a phase file for EDAC.

        @param f: file path or output stream (an object with a write method).

        @param scat: a slice of the self.scattering array belonging to the same atom class.

        @return: None
        """
        if hasattr(f, "write") and callable(f.write):
            energies = np.unique(scat['e'])
            ne = energies.shape[0]
            lmax = scat['l'].max()
            if ne == 1:
                f.write("1 {lmax} regular real\n".format(lmax=lmax))
            else:
                f.write("{nk} E(eV) {lmax} regular real\n".format(nk=ne, lmax=lmax))
            for energy in energies:
                sel = scat['e'] == energy
                idx = np.where(sel)
                energy_scat = scat[idx]
                if ne > 1:
                    f.write("{0:.3f} ".format(energy))
                for item in energy_scat:
                    pha = np.sign(item['t'].real) * np.arcsin(np.sqrt(np.abs(item['t'].imag)))
                    f.write(" {0:.6f}".format(pha))
                for i in range(len(energy_scat), lmax + 1):
                    f.write(" 0")
                f.write("\n")
        else:
            with open(f, "wt", encoding="latin1") as fi:
                self.write_edac_phase_file(fi, scat)

    def parse_radial_file(self, f):
        """
        parse the radial matrix element output file from phagen version 2.2.

        the file contains 7 header lines and one data line per requested energy.
        the data line contains real and imaginary parts of the matrix elements.
        the first four columns contain the electric dipole transitions Rd(li --> li - 1) and Rd(li --> li + 1),
        followed by higher orders that we do not use here.

        @param f: file or path (any file-like or path-like object that can be passed to numpy.genfromtxt).

        @return: None

        @raise ValueError if the file is in a wrong format.
        """
        data = np.atleast_2d(np.genfromtxt(f, skip_header=7))
        if data.shape[0] != self.params.kinetic_energies.shape[0] or data.shape[1] < 4:
            raise ValueError(f"Unexpected array size of Phagen radial matrix elements output: "
                             f"expected ({self.params.kinetic_energies.shape[0]}, >= 4), received {data.shape}")

        self.emission = np.resize(self.emission, data.shape[0:1])
        emission = self.emission
        emission['e'] = self.params.kinetic_energies
        emission['dw'] = data[:, 0] + 1j * data[:, 1]
        emission['up'] = data[:, 2] + 1j * data[:, 3]

    def write_edac_emission_file(self, f):
        """
        write the radial photoemission matrix element in EDAC format.

        requires self.emission, self.params.kinetic_energies and self.params.initial_state.

        @param f: file path or output stream (an object with a write method).

        @return: None
        """
        if hasattr(f, "write") and callable(f.write):
            l0 = self.params.l_init
            energies = self.params.kinetic_energies
            emission = self.emission
            emission['e'] = energies
            ne = energies.shape[0]
            if ne == 1:
                f.write("1 regular2 {l0}\n".format(l0=l0))
            else:
                f.write("{nk} E(eV) regular2 {l0}\n".format(nk=ne, l0=l0))
            for item in emission:
                if ne > 1:
                    f.write("{0:.3f} ".format(item['e']))
                f.write(" {0:.6f} {1:.6f}".format(item['up'].real, item['up'].imag))
                f.write(" {0:.6f} {1:.6f}".format(item['dw'].real, item['dw'].imag))
                f.write("\n")
        else:
            with open(f, "wt", encoding="latin1") as of:
                self.write_edac_emission_file(of)

    def write_edac_emission(self, filename_format):
        """
        write the radial photoemission matrix element in EDAC format.

        requires self.scattering, self.emission, self.params.kinetic_energies and self.params.initial_state.

        @param filename_format: file name including, optionally, a placeholder {} for the atom class.
            since phagen calculates only one emitter, the placeholder is not necessary.

        @return: dictionary that maps atom classes to file names.
            since phagen calculates only one emitter, this dictionary will contain just one entry.
        """
        scat = self.scattering
        atom = scat['a'][0]
        f = filename_format.format(atom)
        self.write_edac_emission_file(f)
        files = {atom: f}
        return files
