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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from pmsco.cluster import Cluster
from pmsco.compat import open

## 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=np.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, "w") 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, "w") 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, "w") as fi:
                self.write_edac_phase_file(fi, scat)

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

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

        @return: None
        """
        dt = [('ar', 'f8'), ('ai', 'f8'), ('br', 'f8'), ('bi', 'f8')]
        data = np.atleast_1d(np.genfromtxt(f, dtype=dt))

        self.emission = np.resize(self.emission, data.shape)
        emission = self.emission
        emission['dw'] = data['ar'] + 1j * data['ai']
        emission['up'] = data['br'] + 1j * data['bi']

    def write_edac_emission(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, "w") as of:
                self.write_edac_emission(of)