# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.


from collections import defaultdict, OrderedDict
import re

import numpy as np

from monty.io import zopen
from monty.json import MSONable

from pymatgen import Orbital, Spin, Element
from pymatgen.electronic_structure.dos import Dos, CompleteDos
from pymatgen.io.feff import Header, Potential, Tags

__author__ = "Alan Dozier, Kiran Mathew, Chen Zheng"
__credits__ = "Anubhav Jain, Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0.3"
__maintainer__ = "Alan Dozier"
__email__ = "adozier@uky.edu"
__status__ = "Beta"
__date__ = "April 7, 2013"

"""
This module defines classes for parsing the FEFF output files.

Currently supports the xmu.dat, ldos.dat output files are for non-spin case.
"""


class LDos(MSONable):
    """
    Parser for ldos files ldos01, ldos02, .....

    Args:
        complete_dos (CompleteDos): complete dos object
        charge_transfer (dict): computed charge transfer between atoms
            dictionary
    """

    def __init__(self, complete_dos, charge_transfer):
        self.complete_dos = complete_dos
        self.charge_transfer = charge_transfer

    @staticmethod
    def from_file(feff_inp_file='feff.inp', ldos_file='ldos'):
        """"
        Creates LDos object from raw Feff ldos files by
        by assuming they are numbered consecutively, i.e. ldos01.dat
        ldos02.dat...

        Args:
            feff_inp_file (str): input file of run to obtain structure
            ldos_file (str): output ldos file of run to obtain dos info, etc.
        """
        header_str = Header.header_string_from_file(feff_inp_file)
        header = Header.from_string(header_str)
        structure = header.struct
        nsites = structure.num_sites
        parameters = Tags.from_file(feff_inp_file)

        if "RECIPROCAL" in parameters:
            pot_dict = dict()
            pot_readstart = re.compile('.*iz.*lmaxsc.*xnatph.*xion.*folp.*')
            pot_readend = re.compile('.*ExternalPot.*switch.*')
            pot_inp = re.sub(r'feff.inp', r'pot.inp', feff_inp_file)
            dos_index = 1
            begin = 0

            with zopen(pot_inp, "r") as potfile:
                for line in potfile:
                    if len(pot_readend.findall(line)) > 0:
                        break

                    if begin == 1:
                        begin += 1
                        continue

                    if begin == 2:
                        z_number = int(line.strip().split()[0])
                        ele_name = Element.from_Z(z_number).name
                        if ele_name not in pot_dict:
                            pot_dict[ele_name] = dos_index
                        else:
                            pot_dict[ele_name] = min(dos_index, pot_dict[ele_name])
                        dos_index += 1

                    if len(pot_readstart.findall(line)) > 0:
                        begin = 1
        else:
            pot_string = Potential.pot_string_from_file(feff_inp_file)
            dicts = Potential.pot_dict_from_string(pot_string)
            pot_dict = dicts[0]

        with zopen(ldos_file + "00.dat", "r") as fobject:
            f = fobject.readlines()
        efermi = float(f[0].split()[4])

        dos_energies = []
        ldos = {}

        for i in range(1, len(pot_dict) + 1):
            if len(str(i)) == 1:
                ldos[i] = np.loadtxt("{}0{}.dat".format(ldos_file, i))
            else:
                ldos[i] = np.loadtxt("{}{}.dat".format(ldos_file, i))

        for i in range(0, len(ldos[1])):
            dos_energies.append(ldos[1][i][0])

        all_pdos = []
        vorb = {"s": Orbital.s, "p": Orbital.py, "d": Orbital.dxy,
                "f": Orbital.f0}
        forb = {"s": 0, "p": 1, "d": 2, "f": 3}

        dlength = len(ldos[1])

        for i in range(nsites):
            pot_index = pot_dict[structure.species[i].symbol]
            all_pdos.append(defaultdict(dict))
            for k, v in vorb.items():
                density = [ldos[pot_index][j][forb[k] + 1]
                           for j in range(dlength)]
                updos = density
                downdos = None
                if downdos:
                    all_pdos[-1][v] = {Spin.up: updos, Spin.down: downdos}
                else:
                    all_pdos[-1][v] = {Spin.up: updos}

        pdos = all_pdos
        vorb2 = {0: Orbital.s, 1: Orbital.py, 2: Orbital.dxy, 3: Orbital.f0}
        pdoss = {structure[i]: {v: pdos[i][v]
                                for v in vorb2.values()}
                 for i in range(len(pdos))}

        forb = {"s": 0, "p": 1, "d": 2, "f": 3}

        tdos = [0] * dlength
        for i in range(nsites):
            pot_index = pot_dict[structure.species[i].symbol]
            for v in forb.values():
                density = [ldos[pot_index][j][v + 1] for j in range(dlength)]
                for j in range(dlength):
                    tdos[j] = tdos[j] + density[j]
        tdos = {Spin.up: tdos}

        dos = Dos(efermi, dos_energies, tdos)
        complete_dos = CompleteDos(structure, dos, pdoss)
        charge_transfer = LDos.charge_transfer_from_file(feff_inp_file,
                                                         ldos_file)
        return LDos(complete_dos, charge_transfer)

    @staticmethod
    def charge_transfer_from_file(feff_inp_file, ldos_file):
        """
        Get charge transfer from file.

        Args:
            feff_inp_file (str): name of feff.inp file for run
            ldos_file (str): ldos filename for run, assume consequetive order,
                i.e., ldos01.dat, ldos02.dat....

        Returns:
            dictionary of dictionaries in order of potential sites
            ({"p": 0.154, "s": 0.078, "d": 0.0, "tot": 0.232}, ...)
        """
        cht = OrderedDict()
        parameters = Tags.from_file(feff_inp_file)

        if 'RECIPROCAL' in parameters:
            dicts = [dict()]
            pot_dict = dict()
            dos_index = 1
            begin = 0
            pot_inp = re.sub(r'feff.inp', r'pot.inp', feff_inp_file)
            pot_readstart = re.compile('.*iz.*lmaxsc.*xnatph.*xion.*folp.*')
            pot_readend = re.compile('.*ExternalPot.*switch.*')
            with zopen(pot_inp, "r") as potfile:
                for line in potfile:
                    if len(pot_readend.findall(line)) > 0:
                        break
                    if begin == 1:
                        z_number = int(line.strip().split()[0])
                        ele_name = Element.from_Z(z_number).name
                        if len(pot_dict) == 0:
                            pot_dict[0] = ele_name
                        elif len(pot_dict) > 0:
                            pot_dict[max(pot_dict.keys()) + 1] = ele_name
                        begin += 1
                        continue
                    if begin == 2:
                        z_number = int(line.strip().split()[0])
                        ele_name = Element.from_Z(z_number).name
                        dicts[0][ele_name] = dos_index
                        dos_index += 1
                        if len(pot_dict) == 0:
                            pot_dict[0] = ele_name
                        elif len(pot_dict) > 0:
                            pot_dict[max(pot_dict.keys()) + 1] = ele_name
                    if len(pot_readstart.findall(line)) > 0:
                        begin = 1
        else:
            pot_string = Potential.pot_string_from_file(feff_inp_file)
            dicts = Potential.pot_dict_from_string(pot_string)
            pot_dict = dicts[1]

        for i in range(0, len(dicts[0]) + 1):
            if len(str(i)) == 1:
                with zopen("{}0{}.dat".format(ldos_file, i), "rt") \
                        as fobject:
                    f = fobject.readlines()
                    s = float(f[3].split()[2])
                    p = float(f[4].split()[2])
                    d = float(f[5].split()[2])
                    f1 = float(f[6].split()[2])
                    tot = float(f[1].split()[4])
                    cht[str(i)] = {pot_dict[i]: {'s': s, 'p': p, 'd': d,
                                                 'f': f1,
                                                 'tot': tot}}
            else:
                with zopen(ldos_file + str(i) + ".dat", "rt") as fid:
                    f = fid.readlines()
                    s = float(f[3].split()[2])
                    p = float(f[4].split()[2])
                    d = float(f[5].split()[2])
                    f1 = float(f[6].split()[2])
                    tot = float(f[1].split()[4])
                    cht[str(i)] = {pot_dict[i]: {'s': s, 'p': p, 'd': d,
                                                 'f': f1,
                                                 'tot': tot}}

        return cht

    def charge_transfer_to_string(self):
        """returns shrage transfer as string"""
        ch = self.charge_transfer
        chts = ['\nCharge Transfer\n\nabsorbing atom']
        for i in range(len(ch)):
            for atom, v2 in ch[str(i)].items():
                a = ['\n', atom, '\n', 's   ', str(v2['s']), '\n',
                     'p   ', str(v2['p']), '\n',
                     'd   ', str(v2['d']), '\n',
                     'f   ', str(v2['f']), '\n',
                     'tot ', str(v2['tot']), '\n']
                chts.extend(a)
        return ''.join(chts)


class Xmu(MSONable):
    """
    Parser for data in 'xmu.dat' file.
    The file 'xmu.dat' contains XANES, EXAFS or NRIXS data depending on the
    situation; \\mu, \\mu_0, and \\chi = \\chi * \\mu_0/ \\mu_0/(edge+50eV) as
    functions of absolute energy E, relative energy E − E_f and wave number k.

    Args:
        header: Header object
        parameters: Tags object
        absorbing_atom (str/int): absorbing atom symbol or index
        data (numpy.ndarray, Nx6): cross_sections

    Default attributes:
        xmu: Photon absorption cross section of absorbing atom in material
        Energies: Energies of data point
        relative_energies: E - E_fermi
        wavenumber: k=\\sqrt(E −E_fermi)
        mu: The total absorption cross-section.
        mu0: The embedded atomic background absorption.
        chi: fine structure.
        Edge: Aborption Edge
        Absorbing atom: Species of absorbing atom
        Material: Formula of material
        Source: Source of structure
        Calculation: Type of Feff calculation performed
    """

    def __init__(self, header, parameters, absorbing_atom, data):
        self.header = header
        self.parameters = parameters
        self.absorbing_atom = absorbing_atom
        self.data = np.array(data)

    @staticmethod
    def from_file(xmu_dat_file="xmu.dat", feff_inp_file="feff.inp"):
        """
        Get Xmu from file.

        Args:
            xmu_dat_file (str): filename and path for xmu.dat
            feff_inp_file (str): filename and path of feff.inp input file

        Returns:
             Xmu object
        """
        data = np.loadtxt(xmu_dat_file)
        header = Header.from_file(feff_inp_file)
        parameters = Tags.from_file(feff_inp_file)
        pots = Potential.pot_string_from_file(feff_inp_file)
        # site index (Note: in feff it starts from 1)
        if "RECIPROCAL" in parameters:
            absorbing_atom = parameters["TARGET"]
        # species symbol
        else:
            absorbing_atom = pots.splitlines()[3].split()[2]
        return Xmu(header, parameters, absorbing_atom, data)

    @property
    def energies(self):
        """
        Returns the absolute energies in eV.
        """
        return self.data[:, 0]

    @property
    def relative_energies(self):
        """
        Returns energy with respect to the fermi level.
        E - E_f
        """
        return self.data[:, 1]

    @property
    def wavenumber(self):
        """
        Returns The wave number in units of \\AA^-1. k=\\sqrt(E −E_f) where E is
        the energy and E_f is the Fermi level computed from electron gas theory
        at the average interstitial charge density.
        """
        return self.data[:, 2]

    @property
    def mu(self):
        """
        Returns the total absorption cross-section.
        """
        return self.data[:, 3]

    @property
    def mu0(self):
        """
        Returns the embedded atomic background absorption.
        """
        return self.data[:, 4]

    @property
    def chi(self):
        """
        Returns the normalized fine structure.
        """
        return self.data[:, 5]

    @property
    def e_fermi(self):
        """
        Returns the fermi level in eV.
        """
        return (self.energies[0] - self.relative_energies[0])

    @property
    def source(self):
        """
        Returns source identification from Header file
        """
        return self.header.source

    @property
    def calc(self):
        """
        Returns type of Feff calculation, XANES or EXAFS
        """
        return "XANES" if "XANES" in self.parameters else "EXAFS"

    @property
    def material_formula(self):
        """
        Returns chemical formula of material from feff.inp file
        """
        try:
            form = self.header.formula
        except IndexError:
            form = 'No formula provided'
        return "".join(map(str, form))

    @property
    def edge(self):
        """
        Returns excitation edge.
        """
        return self.parameters["EDGE"]

    def as_dict(self):
        """
        Returns dict representations of Xmu object
        """
        d = MSONable.as_dict(self)
        d["data"] = self.data.tolist()
        return d


class Eels(MSONable):
    """
    Parse'eels.dat' file.
    """

    def __init__(self, data):
        self.data = np.array(data)

    @property
    def energies(self):
        """
        Returns the energies in eV.
        """
        return self.data[:, 0]

    @property
    def total_spectrum(self):
        """
        Returns the total eels spectrum.
        """
        return self.data[:, 1]

    @property
    def atomic_background(self):
        return self.data[:, 2]

    @property
    def fine_structure(self):
        return self.data[:, 3]

    @staticmethod
    def from_file(eels_dat_file="eels.dat"):
        """
        Parse eels spectrum.

        Args:
            eels_dat_file (str): filename and path for eels.dat

        Returns:
             Eels object
        """
        data = np.loadtxt(eels_dat_file)
        return Eels(data)

    def as_dict(self):
        """
        Returns dict representations of Xmu object
        """
        d = MSONable.as_dict(self)
        d["data"] = self.data.tolist()
        return d
