text: # 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
  