<commit_before>#!/usr/bin/env python
from __future__ import division
"""
Classes for reading/manipulating/writing VASP files.
"""

__author__="<NAME>"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Production"
__date__ ="$Sep 23, 2011M$"

import os
import re
import itertools
import warnings
import xml.sax.handler
import StringIO
from collections import defaultdict
import ConfigParser
import json

import numpy as np
from numpy import identity, array, zeros
from numpy.linalg import det

import pymatgen
from pymatgen.io.io_abc import VaspInput
from pymatgen.util.string_utils import str_aligned, str_delimited
from pymatgen.util.io_utils import file_open_zip_aware, clean_lines, micro_pyawk, clean_json
from pymatgen.core.structure import Structure, Composition
from pymatgen.core.periodic_table import Element
from pymatgen.core.electronic_structure import CompleteDos, Dos, PDos, Spin, Orbital
from pymatgen.core.lattice import Lattice

coord_pattern = re.compile("^\s*([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)")


class VaspParameterSet(object):
    
    def __init__(self, name = "MaterialsProject"):
        self.name = name
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspParameterSets.cfg")))
        self.potcar_settings = dict(self._config.items(name + 'POTCAR'))
        self.incar_settings = dict(self._config.items(name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            self.incar_settings[key] = json.loads(self.incar_settings[key])

class Poscar(VaspInput):
    """
    Object for representing the data in a POSCAR or CONTCAR file.
    """
    
    def __init__(self, struct, comment = None):
        """
        Creates a POSCAR file from a Structure object.
        Arguments:
            struct - Structure object. See pymatgen.core.structure.Structure.
            comment - Optional comment line for POSCAR. Defaults to unit cell formula of structure.
        """

        if struct.is_ordered:
            self._struct = struct
            self._site_symbols = []
            self._natoms = []
            curr_sym = struct[0].specie.symbol
            curr_amt = 0
            for site in struct:
                el = site.specie
                if el.symbol == curr_sym:
                    curr_amt += 1
                else:
                    self._site_symbols.append(curr_sym)
                    self._natoms.append(curr_amt)
                    curr_sym = el.symbol
                    curr_amt = 1
            self._site_symbols.append(curr_sym)
            self._natoms.append(curr_amt)
            self._true_names = True
            self._selective_dynamics = None
            self.comment = struct.formula if comment == None else comment    
        else:
            raise ValueError("Structure with partial occupancies cannot be converted into POSCAR!")

    @property
    def site_symbols(self):
        """
        Symbols for each site in POSCAR.
        """
        return self._site_symbols
    
    @property
    def struct(self):
        """
        Structure associated with the Poscar file.
        """
        return self._struct

    @staticmethod
    def from_file(filename):
        """
        Reads a Poscar from a file.
        The code will try its best to determine the elements in the POSCAR in the following order:
            i) Ideally, if the input file is Vasp5-like and contains element symbols in the 6th line, the code will use that.
            ii) Failing (i), the code will check if a symbol is provided at the end of each coordinate.
            iii) Failing (i) and (ii), the code will try to check if a POTCAR is in the same directory as the POSCAR and use elements from that.
        If all else fails, the code will just assign the first n elements in increasing atomic number, where n is the number of species, 
        to the Poscar.  For example, H, He, Li, ....  This will ensure at least a unique element is assigned to each site and any analysis 
        that does not require specific elemental properties should work fine.
        Arguments:
            filename - file name containing Poscar data.
        Returns:
            Poscar object.
        """
        
        with file_open_zip_aware(filename, "r") as f:
            lines = tuple(clean_lines(f.readlines(), False))

        comment = lines[0]
        scale = float(lines[1])
        lattice = identity(3, float)
        lattice[0] = array([float(s) for s in lines[2].split()])
        lattice[1] = array([float(s) for s in lines[3].split()])
        lattice[2] = array([float(s) for s in lines[4].split()])
        if scale < 0:
            vol = abs(det(lattice))
            lattice = (-scale / vol) ** (1. / 3.) * lattice
        else:
            lattice = scale * lattice
        lattice = Lattice(lattice)

        found_symbols = False
        try:
            natoms = [int(s) for s in lines[5].split()]
        except:
            found_symbols = True

        #Checking for Vasp5+ style Poscars
        if found_symbols:
            symbols = lines[5].split()
            natoms = [int(s) for s in lines[6].split()]
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                atomic_symbols.extend([symbols[i]] * natoms[i])
            ipos = 7
        else:
            ipos = 6

        postype = lines[ipos].split()[0]
        cart = False
        # Selective dynamics
        if postype[0] in 'sS':
            ipos += 1
            postype = lines[ipos].split()[0]
        if postype[0] in 'cCkK':
            cart = True
        N = sum(natoms)

        # See if the element names are appended at the end.
        if not found_symbols:
            try:
                atomic_symbols = [l.split()[3] for l in lines[ipos + 1:ipos + 1 + N]]
                count = 0
                names = list()
                for num in natoms:
                    if not Element.is_valid_symbol(atomic_symbols[count]):
                            raise ValueError("Invalid name")
                    names.append(atomic_symbols[count])
                    count += num
                found_symbols = True
            except:
                pass
            
        #Try to find a element from a POTCAR in the same directory.
        if not found_symbols:
            dirname = os.path.dirname(os.path.abspath(filename))
            for f in os.listdir(dirname):
                if re.search("POTCAR.*",f):
                    try:
                        warnings.warn("POTCAR found! Using elements from POTCAR.")
                        potcar = Potcar.from_file(os.path.join(dirname, f))
                        names = [sym.split("_")[0] for sym in potcar.symbols]
                        atomic_symbols = list()
                        for i in xrange(len(natoms)):
                            atomic_symbols.extend([names[i]] * natoms[i])
                        found_symbols = True
                    except:
                        pass
        
        #Defaulting to false names.
        if not found_symbols:
            names = list()
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                sym = Element.from_Z(i+1).symbol
                names.append(sym)
                atomic_symbols.extend([sym] * natoms[i])
            warnings.warn("Elements in POSCAR cannot be determined. Defaulting to false names, " + " ".join(names)+".")

        # read the atomic coordinates
        coords = zeros((N, 3), float)
        for i in xrange(N):
            iL = ipos + 1 + i
            coords[i] = array([float(s) for s in lines[iL].split()[:3]], float)
            
        struct = Structure(lattice, atomic_symbols, coords, False, False, cart)

        return Poscar(struct,comment)
    
    def get_string(self, direct = True, vasp4_compatible = False):
        """
        Returns a string to be written as a
        POSCAR file. Site symbols are written
        which means compatibilty is for vasp >= 5.
        Optional arguments:
            direct - Whether coordinates are output in direct or cartesian.
            vasp4_compatible - Set to True to omit site symbols on 6th line to maintain backward vasp 4.x compatibility.
        """
        lines = []
        lines += [self.comment]
        lines += ["1.0"]
        lines += [str(self._struct.lattice)]
        if self._true_names == True and not vasp4_compatible:
            lines += [" ".join(self._site_symbols)]
        lines += [" ".join([str(x) for x in self._natoms])]
        if self._selective_dynamics != None:
            lines += ["Selective dynamics"]
            #extra = map(lambda sdrow: " " + ("F", "T")[sdrow[0]] + " " + ("F", "T")[sdrow[1]] + " " + ("F", "T")[sdrow[2]], self._selective_dynamics)
        #else:
        #    extra = None
        if direct == True:
            lines += ["direct"]
            lines += ["%.6f %.6f %.6f %s" % (site.a, site.b, site.c, site.species_string) for site in self._struct.sites]  
        else:
            lines += ["cartesian"]
            lines += ["%.6f %.6f %.6f %s" % (site.x, site.y, site.z, site.species_string) for site in self._struct.sites]

        return "\n".join(lines)    

    def __str__(self):
        """
        String representation of Poscar file.
        """
        return self.get_string()

    def set_site_symbols(self, symbols):
        self._site_symbols = symbols
        self._true_names = True
        #update the Structure as well
        elements = list()
        for i in range(len(symbols)):
            elements.extend([symbols[i]]*self._natoms[i])
        self._struct = Structure(self._struct.lattice, elements, self._struct.frac_coords)

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(str(self) + "\n")
        
VALID_INCAR_TAGS = ("NGX", "NGY", "NGZ", "NGXF", "NGYF", "NGZF", "NBANDS", "NBLK", "SYSTEM", "NWRITE", "ENCUT", "ENAUG",
"PREC", "ISPIN", "MAGMOM", "ISTART", "ICHARG", "INIWAV", "NELM", "NELMIN", "NELMDL", "EDIFF", "EDIFFG", "NSW", "NBLOCK",
"KBLOCK", "IBRION", "NFREE", "POTIM", "ISIF", "PSTRESS", "IWAVPR", "ISYM", "SYMPREC", "LCORR", "TEBEG", "TEEND", "SMASS",
"NPACO", "APACO", "POMASS", "ZVAL", "RWIGS", "LORBIT", "NELECT", "NUPDOWN", "EMIN", "EMAX", "NEDOS", "ISMEAR", "SIGMA", 
"FERWE", "FERDO", "SMEARINGS", "LREAL", "ROPT", "GGA", "VOSKOWN", "LASPH", "ALGO", "IALGO", "LDIAG", "NSIM", "IMIX", "INIMIX", 
"MAXMIX", "AMIX", "BMIX", "AMIX_MAG", "BMIX_MAG", "AMIN", "MIXPRE", "WC", "WEIMIN", "EBREAK", "DEPER", "TIME", "LWAVE", "LCHARG",
"LVTOT", "LELF", "NPAR", "LPLANE","LASYNC", "LSCALAPACK", "LSCALU", "ISPIND", "HFSCREEN", "LHFCALC", "ENCUTFOCK", "NKRED", "LMAXMIX",
"PRECFOCK", "AEXX", "AGGAX", "AGGAC", "ALDAC", "LMAXFOCK", "LMAXFOCKAE", "LTHOMAS", "NKREDX", "NKREDY", "NKREDZ", "EVENONLY", "ODDONLY", "LDAU", "LDAUJ", "LDAUL", "LDAUPRINT", "LDAUTYPE", "LDAUU", "LPEAD", "LCALCPOL", "LCALCEPS", "LEFG", "EFIELD_PEAD", "LNONCOLLINEAR",
"LSORBIT", "IDIPOL", "DIPOL", "LMONO", "LDIPOL", "EPSILON", "EFIELD", "LBERRY", "IGPAR", "NPPSTR", "IPEAD", "I_CONSTRAINED_M", "LAMBDA", "M_CONSTR",
"IMAGES", "SPRING", "LOPTICS", "CSHIFT", "LNABLA", "LEPSILON", "LRPA", "NOMEGA", "NOMEGAR", "LSPECTRAL", "OMEGAMAX", "OMEGATL", "ENCUTGW",
"ENCUTGWSOFT", "ODDONLYGW", "EVENONLYGW", "LSELFENERGY", 'LRHFATM', 'METAGGA', 'LMAXTAU', 'LCOMPAT','ENMAX', 'LMAXPAW', 'LSPIRAL', 'LZEROZ',
'LMETAGGA','ENINI', 'NRMM', 'MREMOVE', 'ADDGRID', 'EFERMI', 'LPARD', 'LSCAAWARE', 'IDIOT', 'LMUSIC', 'LREAL_COMPAT', 'GGA_COMPAT', 'ICORELEVEL', 'LHFONE',
'LRHFCALC', 'LMODELHF', 'ENCUT4O', 'EXXOEP', 'FOURORBIT', 'HFALPHA', 'ALDAX', 'SHIFTRED', 'NMAXFOCKAE', 'HFSCREENC', 'MODEL_GW', 'MODEL_EPS0', 'MODEL_ALPHA',
'LVEL', 'SAXIS', 'QSPIRAL', 'STM', 'KINTER', 'ORBITALMAG', 'LMAGBLOCH', 'LCHIMAG', 'LGAUGE', 'MAGATOM', 'MAGDIPOL', 'AVECCONST', 'LTCTE', 'LTETE',
'L2ORDER', 'LGWLF', 'ENCUTLF', 'LMAXMP2', 'SCISSOR', 'NBANDSGW', 'NBANDSLF', 'DIM', 'ANTIRES', 'LUSEW', 'OMEGAGRID', 'SELFENERGY', 'NKREDLFX', 'NKREDLFY',
'NKREDLFZ', 'MAXMEM', 'TELESCOPE', 'LCRITICAL_MEM', 'GGA2',
'TURBO', 'QUAD_EFG','IRESTART','NREBOOT','NMIN','EREF','KSPACING','KGAMMA','LSUBROT','SCALEE','LVHAR','LORBITALREAL','DARWINR','DARWINV','LFOCKAEDFT','NUCIND','MAGPOS','LNICSALL','LADDER','LHARTREE','IBSE','NBANDSO','NBANDSV','OPTEMAX')

class Incar(dict, VaspInput):
    """
    INCAR object for reading and writing INCAR files
    essentially consists of a dictionary with some helper functions
    """
    
    def __init__(self,params = dict()):
        """
        Creates an Incar object.
        Optional arguments:
            params - A set of input parameters as a dictionary.
        """
        super(Incar,self).__init__()
        self.update(params)

    def __setitem__(self, key, val):
        """
        Add parameter-val pair to Incar.  Warns if parameter is not in list of valid INCAR tags.
        Also cleans the parameter and val by stripping leading and trailing white spaces.
        """
        if key.strip().upper() not in VALID_INCAR_TAGS:
            warnings.warn(key.strip() + " not in VALID_INCAR_TAGS")
        super(Incar,self).__setitem__(key.strip(), Incar.proc_val(key.strip(), val.strip()) if isinstance(val,basestring) else val)

    def get_string(self, sort_keys = False, pretty = False):
        keys = self.keys()
        if sort_keys:
            keys = sorted(keys)
        lines = []
        for k in keys:
            if isinstance(self[k], list):
                lines.append([k," ".join([str(i) for i in self[k]])])
            else:
                lines.append([k,self[k]])
        
        if pretty:
            return str_aligned(lines)
        else:
            return str_delimited(lines, None," = ")
    
    @staticmethod
    def from_structure(structure, parameter_set = VaspParameterSet()):
        incar = Incar()
        poscar = Poscar(structure)
        for key, setting in parameter_set.incar_settings.items():
            if key == "MAGMOM":
                incar[key] =  [setting.get(site.specie.symbol, 0.6) for site in structure]
            elif key in ['LDAUU', 'LDAUJ', 'LDAUL']:
                incar[key] =  [setting.get(sym, 0) for sym in poscar.site_symbols]
            else:
                incar[key] = setting
                
        has_u = sum(incar['LDAUU']) > 0
        if not has_u:
            for key in incar.keys():
                if key.startswith('LDAU'):
                    del incar[key]
        
        return incar
    
    def __str__(self):
        return self.get_string(sort_keys = True, pretty = False)

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename, "r") as f:
            lines = list(clean_lines(f.readlines()))
        params = {}
        for line in lines:
            m = re.match("(\w+)\s*=\s*(.*)", line)
            if m:
                key = m.group(1).strip()
                val = m.group(2).strip()
                val = Incar.proc_val(key, val)
                params[key] = val
        return Incar(params)
    
    @staticmethod
    def proc_val(key, val):
        list_type_keys = ('LDAUU', 'LDAUL', 'LDAUJ', 'LDAUTYPE','MAGMOM')
        boolean_type_keys = ('LDAU', 'LWAVE')
        number_type_keys = ('NSW', 'NELMIN', 'ISIF', 'IBRION', "ISPIN", "EDIFF", "ICHARG", "NELM", "ISMEAR", "NPAR", "SIGMA", "LDAUPRINT", 'LMAXMIX')
         
        
        def smart_int_or_float(numstr):
            if numstr.find(".") != -1:
                return float(numstr)
            else:
                return int(numstr)
        try:
            if key in list_type_keys:
                output = list()
                toks = re.split("\s+", val)
                
                for tok in toks:
                    m = re.match("(\d+)\*([\d\.\-\+]+)", tok)
                    if m:
                        output.extend([smart_int_or_float(m.group(2))] * int(m.group(1)))
                    else:
                        output.append(smart_int_or_float(tok))
                return output
            if key in boolean_type_keys:
                m = re.search("^\W+([TtFf])", val)
                if m:
                    if m.group(1) == "T" or m.group(1) == "t":
                        return True
                    else:
                        return False
                raise ValueError(key + " should be a boolean type!")
            
            if key in number_type_keys:
                return smart_int_or_float(val)
        except:
            return val
              
        return val
            
    def diff(self, other):
        similar_param = {}
        different_param = {}
        for k1,v1 in self.items():
            if k1 not in other:
                different_param[k1] = {"INCAR1": v1, "INCAR2": 'Default'}
            elif v1 != other[k1]:
                different_param[k1] = {"INCAR1": v1, "INCAR2": other[k1]}
            else:
                similar_param[k1] = v1
        for k2,v2 in other.items():
            if k2 not in similar_param and k2 not in different_param:
                if k1 not in self:
                    different_param[k1] = {"INCAR1": 'Default', "INCAR2": v2}
        return {'Same parameters' : similar_param, 'Different': different_param}
                
    def __add__(self, other):
        """
        Add all the values of another INCAR object to this object
        Facilitates the use of "standard" INCARs
        """
        params = {k:v for k,v in self.items()}
        for k, v in other.items():
            if k in self and v != self[k]:
                raise ValueError("Incars have conflicting values!")
            else:
                params[k] = v
        return Incar(params)

class Kpoints(VaspInput):
    """
    Very basic KPOINT reader/writer
    The lines of the original file
    are stored as is but the kpts are stored
    as a list which can be accessed and
    modified
    """
    def __init__(self, comment = "Default", num_kpts = 0, style = "Gamma", kpts = [[1,1,1]], kpts_shift = [0,0,0]):
        self.comment = comment
        self.num_kpts = num_kpts
        self.style = style
        self.kpts = kpts
        self.kpts_shift = kpts_shift

    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename) as f:
            lines = list(clean_lines(f.readlines(), False))
        kpoints = Kpoints()
        kpoints.comment = lines[0].strip()
        kpoints.num_kpts = int(lines[1].strip())
        style = lines[2].strip().lower()[0]
        if style == "a":
            kpoints.style = "Auto"
            kpoints.kpts = [[int(lines[3].strip())]]
        elif style == "g" or style == "m":
            kpoints.style = "Gamma" if style == "g" else "Monkhorst-Pack"
            kpoints.kpts = [[int(x) for x in lines[3].strip().split()]]
            if len(lines)>4 and coord_pattern.match(lines[4].strip()):
                try:
                    kpoints.kpts_shift = [int(x) for x in lines[4].strip().split()]
                except:
                    pass
        elif kpoints.num_kpts <= 0:
            kpoints.style = "Cartesian" if style == "c" else "Reciprocal"
            kpoints.kpts = [[float(x) for x in lines[i].strip().split()] for i in xrange(3,6)]
            kpoints.kpts_shift = [float(x) for x in lines[6].strip().split()]
                
            
        return kpoints

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    def __str__(self):
        lines = []
        lines += [self.comment]
        lines += [str(self.num_kpts)]
        lines += [self.style]
        for line in self.kpts:
            lines += [" ".join([str(x) for x in line])]
        lines += [" ".join([str(x) for x in self.kpts_shift])]
        return "\n".join(lines)
    
    @property
    def to_dict(self):
        d = {'comment': self.comment, 'nkpoints' : self.num_kpts, 'generation_style' : self.style, 'kpoints': self.kpts, 'usershift': self.kpts_shift}
        optional_paras = ['genvec1', 'genvec2', 'genvec3', 'shift']
        for para in optional_paras:
            if para in self.__dict__:
                d[para] = self.__dict__[para]
        return d 
    
class PotcarSingle(VaspInput):
    """
    Object for a **single** POTCAR.
    The builder assumes the complete string is the POTCAR
    contains the complete untouched data in "data" as a string and
    a dictionary of keywords
    """
    def __init__(self, data):
        """
        Expects a complete and single potcar file as a string in "data"
        """
        self.data = data # raw POTCAR as a string
        keyValPairs = re.compile(r";*\s*(.+?)\s*=\s*([^;\n]+)\s*", re.M).findall(data)
        self.keywords = dict(keyValPairs) # all key = val found in the POTCAR as a dictionary all keys and vals are strings

    def __str__(self):
        return self.data

    def write_file(self, filename):
        writer = open(filename, 'w')
        writer.write(self.__str__() + "\n")
        writer.close()

    @property
    def symbol(self):
        """
        Full name of POTCAR, e.g., Fe_pv
        """
        return self.keywords['TITEL'].split(" ")[1].strip()        
    
    @property
    def element(self):
        """
        Attempt to return the atomic symbol based on the VRHFIN keyword
        """
        return self.keywords['VRHFIN'].split(":")[0].strip()

    @property
    def atomic_no(self):
        """
        Attempt to return the atomic number based on the VRHFIN keyword
        """
        return Element(self.element).Z

class Potcar(list,VaspInput):
    """
    Object for reading and writing POTCAR files for
    calculations.
    """
    functional_dir = {'PBE':'POT_GGA_PAW_PBE', 'LDA':'POT_LDA_PAW', 'PW91':'POT_GGA_PAW_PW91'}
    
    def __init__(self, symbols = None):
        if symbols != None:
            self.set_symbols(symbols)

    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename, "r") as reader:
            fData = reader.read()
        potcar = Potcar()
        potcar_strings = re.compile(r"\n{0,1}\s*(.*?End of Dataset)", re.S).findall(fData)
        for p in potcar_strings:
            potcar.append(PotcarSingle(p))
        return potcar

    def __str__(self):
        return "\n".join([str(potcar) for potcar in self])

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @property
    def symbols(self):
        """
        Get the atomic symbols of all the atoms in the POTCAR file
        """
        return [p.symbol for p in self]

    def set_symbols(self, elements, functional = 'PBE', parameter_set = VaspParameterSet()):
        module_dir = os.path.dirname(pymatgen.__file__)
        config = ConfigParser.SafeConfigParser()
        config.readfp(open(os.path.join(module_dir, "pymatgen.cfg")))
        VASP_PSP_DIR = os.path.join(config.get('VASP', 'pspdir'), Potcar.functional_dir[functional])
        
        del self[:]
        for el in elements:
            sym = el
            if parameter_set and Element.is_valid_symbol(el) and el in parameter_set.potcar_settings:
                sym = parameter_set.potcar_settings[el]
            reader = file_open_zip_aware(os.path.join(VASP_PSP_DIR, "POTCAR." + sym + ".gz"))
            self.append(PotcarSingle(reader.read()))
            reader.close()

class Vasprun(object):
    """
    Vastly improved sax-based parser for vasprun.xml files.
    Speedup over Dom is at least 2x for smallish files (~1Mb) to orders of magnitude for larger files (~10Mb).
    All data is stored as attributes, which are delegated to the VasprunHandler object.
    Accessible attributes from VasprunHandler are:
        Vasp results
        ------------
        energies - All energies in run, represented as a list of ionic steps, with a list of scstep energies.
                   E.g. [ [0.1, 0.2, 0.1], [0.1], [0.22, 0.23, 0.20, 0.205]] 
        structures - List of Structure objects for the structure at each ionic step.
        tdos - Total dos calculated at the end of run.
        idos - Integrated dos calculated at the end of run.
        pdos - List of list of PDos objects. Access as pdos[atomindex][orbitalindex]
        efermi - Fermi energy
        eigenvalues - Final eigenvalues as a dict of {(kpoint index, Spin.up):[[eigenvalue, occu]]}. 
                      This representation is probably not ideal, but since this is not used anywhere else for now, I leave it as such.
                      Future developers who need to work with this should refactored the object into a sensible structure.
        
        Vasp inputs
        -----------
        incar - Incar object for parameters specified in INCAR file.
        parameters - Incar object with parameters that vasp actually used, including all defaults.
        kpoints = Kpoints object for KPOINTS specified in run.
        actual_kpoints - List of actual kpoints, e.g., [[0.25, 0.125, 0.08333333], [-0.25, 0.125, 0.08333333], [0.25, 0.375, 0.08333333], ....]
        actual_kpoints_weights = List of kpoint weights, E.g., [0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, ....]
        atomic_symbols - List of atomic symbols, e.g., [u'Li', u'Fe', u'Fe', u'P', u'P', u'P']
        potcar_symbols - List of POTCAR symbols. E.g., [u'PAW_PBE Li 17Jan2003', u'PAW_PBE Fe 06Sep2000', ..]
    
    A few helper attributes have also been added to get commonly used results such as final energies.
    
    Author: <NAME>
    """
    supported_properties = ['vasp_version', 'incar', 'parameters', 'potcar_symbols', 'atomic_symbols', 'kpoints', 'actual_kpoints', 'structures',
                            'actual_kpoints_weights', 'dos_energies', 'eigenvalues', 'tdos', 'idos', 'pdos', 'efermi', 'ionic_steps', 'dos_has_errors']
    
    def __init__(self, filename):
        self._filename = filename    
        with file_open_zip_aware(filename) as f:
            self._handler = VasprunHandler(filename)
            self._parser = xml.sax.parse(f, self._handler)
            for k in Vasprun.supported_properties:
                setattr(self, k, getattr(self._handler, k))
    
    @property
    def converged(self):
        return len(self.structures) - 2 < self.parameters['NSW'] or self.parameters['NSW'] == 0
    
    @property
    def final_energy(self):
        """
        Final energy from the vasp run.
        """
        return self.ionic_steps[-1]['electronic_steps'][-1]['e_wo_entrp']
    
    @property
    def final_structure(self):
        """
        Final structure from vasprun.
        """
        return self.structures[-1]
    
    @property
    def initial_structure(self):
        """
        Initial structure from vasprun.
        """
        return self.structures[0]
    
    @property
    def complete_dos(self):
        """
        A complete dos object which incorporates the total dos and all projected dos.
        """
        return CompleteDos(self.final_structure, self.tdos, self.pdos)
    
    @property
    def to_dict(self):
        d = {}
        d['vasp_version'] = self.vasp_version
        d['has_vasp_completed'] = self.converged
        d['nsites'] = len(self.final_structure)
        d['unit_cell_formula'] = self.final_structure.composition.to_dict
        comp = self.final_structure.composition
        d['reduced_cell_formula'] = Composition.from_formula(comp.reduced_formula).to_dict
        d['pretty_formula'] = comp.reduced_formula
        symbols = [re.split("\s+", s)[1] for s in self.potcar_symbols]
        symbols = [re.split("_", s)[0] for s in symbols]
        d['elements'] = symbols
        d['nelements'] = len(symbols)
        d['is_hubbard'] = self.incar.get('LDAU', False)
        if d['is_hubbard']:
            us = self.incar.get('LDAUU', self.parameters.get('LDAUU'))
            js = self.incar.get('LDAUJ', self.parameters.get('LDAUJ'))
            if len(us) == len(symbols):
                d['hubbards'] = { symbols[i] : us[i] - js[i] for i in xrange(len(symbols))}
            elif sum(us) == 0 and sum(js) == 0:
                d['is_hubbard'] = False
                d['hubbards'] = {}
            else:
                raise VaspParserError("Length of U value parameters and atomic symbols are mismatched")
        else:
            d['hubbards'] = {}
        if d['is_hubbard']:
            d['run_type'] = "GGA+U"
        elif self.parameters.get('LHFCALC', False):
            d['run_type'] = "HF"
        else:
            d['run_type'] = "GGA"
                   
        d['input'] = {}
        d['input']['incar'] = {k:v for k,v in self.incar.items()}
        d['input']['crystal'] = self.initial_structure.to_dict
        d['input']['kpoints'] = self.kpoints.to_dict
        d['input']['kpoints']['actual_points'] = [{'abc':list(self.actual_kpoints[i]), 'weight':self.actual_kpoints_weights[i]} for i in xrange(len(self.actual_kpoints))] 
        d['input']['potcar'] = [s.split(" ")[1] for s in self.potcar_symbols]
        d['input']['parameters'] = {k:v for k,v in self.parameters.items()}
        
        d['output'] = {}
        d['output']['ionic_steps'] = clean_json(self.ionic_steps)
        d['output']['final_energy'] = self.final_energy
        d['output']['final_energy_per_atom'] = self.final_energy / len(self.final_structure)
        d['output']['crystal'] = self.final_structure.to_dict
        d['output']['efermi'] = self.efermi
        #{(kpoint index, Spin.up):array(float)}
        
        d['output']['eigenvalues'] = {}
        for (index, spin), values in self.eigenvalues.items():
            if str(index) not in d['output']['eigenvalues']:
                d['output']['eigenvalues'][str(index)] = {str(spin):values}
            else:
                d['output']['eigenvalues'][str(index)][str(spin)] = values
                
        return d

 
class VasprunHandler(xml.sax.handler.ContentHandler):
    """
    Sax handler for vasprun.xml.
    Attributes are mirrored into Vasprun object.
    
    Author: <NAME>
    """
        
    def __init__(self, filename):
        self.filename = filename
        # variables to be filled
        self.vasp_version = None
        self.incar = Incar()
        self.parameters = Incar()
        self.potcar_symbols = []
        self.atomic_symbols = []
        self.kpoints = Kpoints()
        self.actual_kpoints = []
        self.actual_kpoints_weights = []
        self.dos_energies = None
        self.eigenvalues = {}#  will  be  {(kpoint index, Spin.up):array(float)}
        self.tdos = {}
        self.idos = {}
        self.pdos = {}
        self.efermi = None 
        self.ionic_steps = [] # should be a list of dict
        self.structures = []
        
        self.input_read = False
        self.all_calculations_read = False
        self.read_structure = False
        self.read_calculation = False
        self.read_eigen = False
        self.read_dos= False
        self.in_efermi = False
        self.read_atoms = False
        self.read_lattice = False
        self.read_positions = False
        self.incar_param = None
        self.dos_energies_val = []
        self.dos_val = []
        self.idos_val = []
        self.raw_data = []
        self.dos_has_errors = False #will be set to true if there is an error parsing the Dos.
        self.state = defaultdict(bool)
        
    
    def in_all(self, xml_tags):
        return all([getattr(self, 'in_'+tag, None) for tag in xml_tags])
    
    def startElement(self, name, attributes):
        
        self.state[name] = True if 'name' not in attributes else attributes['name']
        self.read_val = False
        
        #Nested if loops makes reading much faster.
        if not self.input_read: #reading input parameters
            if (name == "i" or name == "v") and (self.state['incar'] or self.state['parameters']):
                self.incar_param = attributes['name']
                self.param_type = 'float' if 'type' not in attributes else attributes['type']
                self.read_val = True
            elif name == "v" and self.state['kpoints']:
                self.read_val = True
            elif name == "generation" and self.state['kpoints']:
                self.kpoints.comment   = "Kpoints from vasprun.xml"
                self.kpoints.num_kpts  = 0
                self.kpoints.style = attributes['param']
                self.kpoints.kpts = []
                self.kpoints.kpts_shift = [0,0,0]
            elif name == "c" and (self.state['array'] == "atoms" or self.state['array'] == "atomtypes"):
                self.read_val = True
            elif name == "i" and self.state['i'] == "version" and self.state['generator']:
                self.read_val = True
                
        else: #reading calculations and structures.
            if self.read_calculation:
                if name == "i" and self.state['scstep']:
                    self.read_val = True
                elif name == "v" and (self.state['varray'] == "forces" or self.state['varray'] == "stress"):
                    self.read_positions = True
            if self.read_structure:
                if name == "v" and self.state['varray'] == 'basis':
                    self.read_lattice = True
                elif name == "v" and self.state['varray'] == 'positions':
                    self.read_positions = True
            if name == "calculation":
                self.scdata = []
                self.read_calculation = True
            elif name == "scstep":
                self.scstep = {}
            elif name == 'structure':
                self.latticestr = StringIO.StringIO()
                self.posstr = StringIO.StringIO()
                self.read_structure = True
            elif name == 'varray' and (self.state['varray'] == "forces" or self.state['varray'] == "stress"):
                self.posstr = StringIO.StringIO()
            elif name == "eigenvalues":
                self.all_calculations_read = True
            if self.read_eigen:
                if name == "r" and self.state["set"]:
                    self.read_val = True
                elif name == "set" and "comment" in attributes:
                    comment = attributes["comment"]
                    self.state["set"] = comment
                    if comment.startswith("spin"):
                        self.eigen_spin = Spin.up if self.state["set"] == "spin 1" else Spin.down
                    if comment.startswith("kpoint"):
                        self.eigen_kpoint = int(comment.split(" ")[1])
            elif self.read_dos:
                if (name == "i" and self.state["i"] == "efermi") or (name == "r" and self.state["set"] ):
                    self.read_val = True
                elif name == "set" and "comment" in attributes:
                    comment = attributes["comment"]
                    self.state["set"] = comment
                    if self.state['partial']:
                        if comment.startswith("ion"):
                            self.pdos_ion = int(comment.split(" ")[1])
                        elif comment.startswith("spin"):
                            self.pdos_spin = Spin.up if self.state["set"] == "spin 1" else Spin.down
            elif name == "dos":
                self.dos_energies = None
                self.tdos = {}
                self.idos = {}
                self.pdos = {}
                self.efermi = None 
                self.read_dos = True
            elif name == "eigenvalues":
                self.eigenvalues = {}#  will  be  {(kpoint index, Spin.up):array(float)}
                self.read_eigen = True
            
        if self.read_val:
            self.val = StringIO.StringIO()

    def characters(self, data):
        if self.read_val:
            self.val.write(data)
        if self.read_lattice:
            self.latticestr.write(data)
        elif self.read_positions:
            self.posstr.write(data)
    
    #To correct for stupid vasp bug which names Xenon as X!!
    EL_MAPPINGS = {'X':'Xe'}
    
    def endElement(self, name):
        
        if not self.input_read:
            if name == "i":
                if self.state['incar']:
                    self.incar[self.incar_param] = parse_parameters(self.param_type, self.val.getvalue().strip())
                elif self.state['parameters']:
                    self.parameters[self.incar_param] = parse_parameters(self.param_type, self.val.getvalue().strip())
                elif self.state['generator'] and self.state["i"] == "version":
                    self.vasp_version = self.val.getvalue().strip()
                self.incar_param = None
            elif name == "set":
                if self.state['array'] == "atoms":
                    self.atomic_symbols = self.atomic_symbols[::2]
                    self.atomic_symbols = [sym if sym not in VasprunHandler.EL_MAPPINGS else VasprunHandler.EL_MAPPINGS[sym] for sym in self.atomic_symbols]
                elif self.state['array'] == "atomtypes":
                    self.potcar_symbols = self.potcar_symbols[4::5]
                    self.input_read = True
            elif name == "c":
                if self.state['array'] == "atoms":
                    self.atomic_symbols.append(self.val.getvalue().strip())
                elif self.state['array'] == "atomtypes":
                    self.potcar_symbols.append(self.val.getvalue().strip())
            elif name == "v":
                if self.state['incar']:
                    self.incar[self.incar_param] = parse_v_parameters(self.param_type, self.val.getvalue().strip(), self.filename, self.incar_param)
                    self.incar_param = None
                elif self.state['parameters']:
                    self.parameters[self.incar_param] = parse_v_parameters(self.param_type, self.val.getvalue().strip(), self.filename, self.incar_param)
                elif self.state['kpoints']:
                    if self.state['varray'] == 'kpointlist':
                        self.actual_kpoints.append([float(x) for x in re.split("\s+",self.val.getvalue().strip())])
                    if self.state['varray'] == 'weights':
                        self.actual_kpoints_weights.append(float(self.val.getvalue()))
                    if self.state['v'] == "divisions":
                        self.kpoints.kpts   = [[int(x) for x in re.split("\s+",self.val.getvalue().strip())]]
                    elif self.state['v'] == "usershift":
                        self.kpoints.kpts_shift   = [float(x) for x in re.split("\s+",self.val.getvalue().strip())]
                    elif self.state['v'] == "genvec1" or self.state['v'] == "genvec2" or self.state['v'] == "genvec3" or self.state['v'] == "shift":
                        setattr(self.kpoints, self.state['v'], [float(x) for x in re.split("\s+",self.val.getvalue().strip())])
                        
        else:
            if self.read_calculation:
                if name == "i" and self.state['scstep']:
                    self.scstep[self.state['i']] = float(self.val.getvalue())
                elif name == 'scstep':
                    self.scdata.append(self.scstep)
                elif name == 'varray' and self.state['varray'] == "forces":
                    self.forces = np.array([float(x) for x in re.split("\s+",self.posstr.getvalue().strip())])
                    self.forces.shape = (len(self.atomic_symbols), 3)
                elif name == 'varray' and self.state['varray'] == "stress":
                    self.stress = np.array([float(x) for x in re.split("\s+",self.posstr.getvalue().strip())])
                    self.stress.shape = (3, 3)
                elif name == "calculation":
                    self.ionic_steps.append({'electronic_steps':self.scdata, 'structure':self.structures[-1], 'forces': self.forces, 'stress':self.stress})
                    self.read_calculation = False
            if self.read_structure:
                if name == "v":
                    self.read_positions = False
                    self.read_lattice = False
                elif name == "structure":
                    self.lattice = np.array([float(x) for x in re.split("\s+",self.latticestr.getvalue().strip())])
                    self.lattice.shape = (3,3)
                    self.pos = np.array([float(x) for x in re.split("\s+",self.posstr.getvalue().strip())])
                    self.pos.shape = (len(self.atomic_symbols), 3)
                    self.structures.append(Structure(self.lattice, self.atomic_symbols, self.pos))
                    self.read_structure = False
            elif self.read_dos:
                try:
                    if name == "i" and self.state["i"] == "efermi":
                        self.efermi = float(self.val.getvalue().strip())
                    elif name == "r" and self.state["total"]  and str(self.state["set"]).startswith("spin"):
                        tok = re.split("\s+", self.val.getvalue().strip())
                        self.dos_energies_val.append(float(tok[0]))
                        self.dos_val.append(float(tok[1]))
                        self.idos_val.append(float(tok[2]))
                    elif name == "r" and self.state["partial"]  and str(self.state["set"]).startswith("spin"):
                        tok = re.split("\s+", self.val.getvalue().strip())
                        self.raw_data.append([float(i) for i in tok[1:]])
                    elif name == "set" and self.state["total"] and str(self.state["set"]).startswith("spin"):
                        spin = Spin.up if self.state["set"] == "spin 1" else Spin.down
                        self.tdos[spin] = self.dos_val
                        self.idos[spin] = self.dos_val
                        self.dos_energies = self.dos_energies_val
                        self.dos_energies_val = []
                        self.dos_val = []
                        self.idos_val = []
                    elif name == "set" and self.state["partial"] and str(self.state["set"]).startswith("spin"):
                        spin = Spin.up if self.state["set"] == "spin 1" else Spin.down<commit_msg>add the possibility to get the reciprocal lattice from vasprun.xml<commit_after>#!/usr/bin/env python
from __future__ import division
"""
Classes for reading/manipulating/writing VASP files.
"""

__author__="<NAME>"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Production"
__date__ ="$Sep 23, 2011M$"

import os
import re
import itertools
import warnings
import xml.sax.handler
import StringIO
from collections import defaultdict
import ConfigParser
import json

import numpy as np
from numpy import identity, array, zeros
from numpy.linalg import det

import pymatgen
from pymatgen.io.io_abc import VaspInput
from pymatgen.util.string_utils import str_aligned, str_delimited
from pymatgen.util.io_utils import file_open_zip_aware, clean_lines, micro_pyawk, clean_json
from pymatgen.core.structure import Structure, Composition
from pymatgen.core.periodic_table import Element
from pymatgen.core.electronic_structure import CompleteDos, Dos, PDos, Spin, Orbital
from pymatgen.core.lattice import Lattice

coord_pattern = re.compile("^\s*([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)\s+([\d+\.\-Ee]+)")


class VaspParameterSet(object):
    
    def __init__(self, name = "MaterialsProject"):
        self.name = name
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspParameterSets.cfg")))
        self.potcar_settings = dict(self._config.items(name + 'POTCAR'))
        self.incar_settings = dict(self._config.items(name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            self.incar_settings[key] = json.loads(self.incar_settings[key])
        

class Poscar(VaspInput):
    """
    Object for representing the data in a POSCAR or CONTCAR file.
    """
    
    def __init__(self, struct, comment = None):
        """
        Creates a POSCAR file from a Structure object.
        Arguments:
            struct - Structure object. See pymatgen.core.structure.Structure.
            comment - Optional comment line for POSCAR. Defaults to unit cell formula of structure.
        """

        if struct.is_ordered:
            self._struct = struct
            self._site_symbols = []
            self._natoms = []
            curr_sym = struct[0].specie.symbol
            curr_amt = 0
            for site in struct:
                el = site.specie
                if el.symbol == curr_sym:
                    curr_amt += 1
                else:
                    self._site_symbols.append(curr_sym)
                    self._natoms.append(curr_amt)
                    curr_sym = el.symbol
                    curr_amt = 1
            self._site_symbols.append(curr_sym)
            self._natoms.append(curr_amt)
            self._true_names = True
            self._selective_dynamics = None
            self.comment = struct.formula if comment == None else comment    
        else:
            raise ValueError("Structure with partial occupancies cannot be converted into POSCAR!")

    @property
    def site_symbols(self):
        """
        Symbols for each site in POSCAR.
        """
        return self._site_symbols
    
    @property
    def struct(self):
        """
        Structure associated with the Poscar file.
        """
        return self._struct

    @staticmethod
    def from_file(filename):
        """
        Reads a Poscar from a file.
        The code will try its best to determine the elements in the POSCAR in the following order:
            i) Ideally, if the input file is Vasp5-like and contains element symbols in the 6th line, the code will use that.
            ii) Failing (i), the code will check if a symbol is provided at the end of each coordinate.
            iii) Failing (i) and (ii), the code will try to check if a POTCAR is in the same directory as the POSCAR and use elements from that.
        If all else fails, the code will just assign the first n elements in increasing atomic number, where n is the number of species, 
        to the Poscar.  For example, H, He, Li, ....  This will ensure at least a unique element is assigned to each site and any analysis 
        that does not require specific elemental properties should work fine.
        Arguments:
            filename - file name containing Poscar data.
        Returns:
            Poscar object.
        """
        
        with file_open_zip_aware(filename, "r") as f:
            lines = tuple(clean_lines(f.readlines(), False))

        comment = lines[0]
        scale = float(lines[1])
        lattice = identity(3, float)
        lattice[0] = array([float(s) for s in lines[2].split()])
        lattice[1] = array([float(s) for s in lines[3].split()])
        lattice[2] = array([float(s) for s in lines[4].split()])
        if scale < 0:
            vol = abs(det(lattice))
            lattice = (-scale / vol) ** (1. / 3.) * lattice
        else:
            lattice = scale * lattice
        lattice = Lattice(lattice)

        found_symbols = False
        try:
            natoms = [int(s) for s in lines[5].split()]
        except:
            found_symbols = True

        #Checking for Vasp5+ style Poscars
        if found_symbols:
            symbols = lines[5].split()
            natoms = [int(s) for s in lines[6].split()]
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                atomic_symbols.extend([symbols[i]] * natoms[i])
            ipos = 7
        else:
            ipos = 6

        postype = lines[ipos].split()[0]
        cart = False
        # Selective dynamics
        if postype[0] in 'sS':
            ipos += 1
            postype = lines[ipos].split()[0]
        if postype[0] in 'cCkK':
            cart = True
        N = sum(natoms)

        # See if the element names are appended at the end.
        if not found_symbols:
            try:
                atomic_symbols = [l.split()[3] for l in lines[ipos + 1:ipos + 1 + N]]
                count = 0
                names = list()
                for num in natoms:
                    if not Element.is_valid_symbol(atomic_symbols[count]):
                            raise ValueError("Invalid name")
                    names.append(atomic_symbols[count])
                    count += num
                found_symbols = True
            except:
                pass
            
        #Try to find a element from a POTCAR in the same directory.
        if not found_symbols:
            dirname = os.path.dirname(os.path.abspath(filename))
            for f in os.listdir(dirname):
                if re.search("POTCAR.*",f):
                    try:
                        warnings.warn("POTCAR found! Using elements from POTCAR.")
                        potcar = Potcar.from_file(os.path.join(dirname, f))
                        names = [sym.split("_")[0] for sym in potcar.symbols]
                        atomic_symbols = list()
                        for i in xrange(len(natoms)):
                            atomic_symbols.extend([names[i]] * natoms[i])
                        found_symbols = True
                    except:
                        pass
        
        #Defaulting to false names.
        if not found_symbols:
            names = list()
            atomic_symbols = list()
            for i in xrange(len(natoms)):
                sym = Element.from_Z(i+1).symbol
                names.append(sym)
                atomic_symbols.extend([sym] * natoms[i])
            warnings.warn("Elements in POSCAR cannot be determined. Defaulting to false names, " + " ".join(names)+".")

        # read the atomic coordinates
        coords = zeros((N, 3), float)
        for i in xrange(N):
            iL = ipos + 1 + i
            coords[i] = array([float(s) for s in lines[iL].split()[:3]], float)
            
        struct = Structure(lattice, atomic_symbols, coords, False, False, cart)

        return Poscar(struct,comment)
    
    def get_string(self, direct = True, vasp4_compatible = False):
        """
        Returns a string to be written as a
        POSCAR file. Site symbols are written
        which means compatibilty is for vasp >= 5.
        Optional arguments:
            direct - Whether coordinates are output in direct or cartesian.
            vasp4_compatible - Set to True to omit site symbols on 6th line to maintain backward vasp 4.x compatibility.
        """
        lines = []
        lines += [self.comment]
        lines += ["1.0"]
        lines += [str(self._struct.lattice)]
        if self._true_names == True and not vasp4_compatible:
            lines += [" ".join(self._site_symbols)]
        lines += [" ".join([str(x) for x in self._natoms])]
        if self._selective_dynamics != None:
            lines += ["Selective dynamics"]
            #extra = map(lambda sdrow: " " + ("F", "T")[sdrow[0]] + " " + ("F", "T")[sdrow[1]] + " " + ("F", "T")[sdrow[2]], self._selective_dynamics)
        #else:
        #    extra = None
        if direct == True:
            lines += ["direct"]
            lines += ["%.6f %.6f %.6f %s" % (site.a, site.b, site.c, site.species_string) for site in self._struct.sites]  
        else:
            lines += ["cartesian"]
            lines += ["%.6f %.6f %.6f %s" % (site.x, site.y, site.z, site.species_string) for site in self._struct.sites]

        return "\n".join(lines)    

    def __str__(self):
        """
        String representation of Poscar file.
        """
        return self.get_string()

    def set_site_symbols(self, symbols):
        self._site_symbols = symbols
        self._true_names = True
        #update the Structure as well
        elements = list()
        for i in range(len(symbols)):
            elements.extend([symbols[i]]*self._natoms[i])
        self._struct = Structure(self._struct.lattice, elements, self._struct.frac_coords)

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(str(self) + "\n")
        
VALID_INCAR_TAGS = ("NGX", "NGY", "NGZ", "NGXF", "NGYF", "NGZF", "NBANDS", "NBLK", "SYSTEM", "NWRITE", "ENCUT", "ENAUG",
"PREC", "ISPIN", "MAGMOM", "ISTART", "ICHARG", "INIWAV", "NELM", "NELMIN", "NELMDL", "EDIFF", "EDIFFG", "NSW", "NBLOCK",
"KBLOCK", "IBRION", "NFREE", "POTIM", "ISIF", "PSTRESS", "IWAVPR", "ISYM", "SYMPREC", "LCORR", "TEBEG", "TEEND", "SMASS",
"NPACO", "APACO", "POMASS", "ZVAL", "RWIGS", "LORBIT", "NELECT", "NUPDOWN", "EMIN", "EMAX", "NEDOS", "ISMEAR", "SIGMA", 
"FERWE", "FERDO", "SMEARINGS", "LREAL", "ROPT", "GGA", "VOSKOWN", "LASPH", "ALGO", "IALGO", "LDIAG", "NSIM", "IMIX", "INIMIX", 
"MAXMIX", "AMIX", "BMIX", "AMIX_MAG", "BMIX_MAG", "AMIN", "MIXPRE", "WC", "WEIMIN", "EBREAK", "DEPER", "TIME", "LWAVE", "LCHARG",
"LVTOT", "LELF", "NPAR", "LPLANE","LASYNC", "LSCALAPACK", "LSCALU", "ISPIND", "HFSCREEN", "LHFCALC", "ENCUTFOCK", "NKRED", "LMAXMIX",
"PRECFOCK", "AEXX", "AGGAX", "AGGAC", "ALDAC", "LMAXFOCK", "LMAXFOCKAE", "LTHOMAS", "NKREDX", "NKREDY", "NKREDZ", "EVENONLY", "ODDONLY", "LDAU", "LDAUJ", "LDAUL", "LDAUPRINT", "LDAUTYPE", "LDAUU", "LPEAD", "LCALCPOL", "LCALCEPS", "LEFG", "EFIELD_PEAD", "LNONCOLLINEAR",
"LSORBIT", "IDIPOL", "DIPOL", "LMONO", "LDIPOL", "EPSILON", "EFIELD", "LBERRY", "IGPAR", "NPPSTR", "IPEAD", "I_CONSTRAINED_M", "LAMBDA", "M_CONSTR",
"IMAGES", "SPRING", "LOPTICS", "CSHIFT", "LNABLA", "LEPSILON", "LRPA", "NOMEGA", "NOMEGAR", "LSPECTRAL", "OMEGAMAX", "OMEGATL", "ENCUTGW",
"ENCUTGWSOFT", "ODDONLYGW", "EVENONLYGW", "LSELFENERGY", 'LRHFATM', 'METAGGA', 'LMAXTAU', 'LCOMPAT','ENMAX', 'LMAXPAW', 'LSPIRAL', 'LZEROZ',
'LMETAGGA','ENINI', 'NRMM', 'MREMOVE', 'ADDGRID', 'EFERMI', 'LPARD', 'LSCAAWARE', 'IDIOT', 'LMUSIC', 'LREAL_COMPAT', 'GGA_COMPAT', 'ICORELEVEL', 'LHFONE',
'LRHFCALC', 'LMODELHF', 'ENCUT4O', 'EXXOEP', 'FOURORBIT', 'HFALPHA', 'ALDAX', 'SHIFTRED', 'NMAXFOCKAE', 'HFSCREENC', 'MODEL_GW', 'MODEL_EPS0', 'MODEL_ALPHA',
'LVEL', 'SAXIS', 'QSPIRAL', 'STM', 'KINTER', 'ORBITALMAG', 'LMAGBLOCH', 'LCHIMAG', 'LGAUGE', 'MAGATOM', 'MAGDIPOL', 'AVECCONST', 'LTCTE', 'LTETE',
'L2ORDER', 'LGWLF', 'ENCUTLF', 'LMAXMP2', 'SCISSOR', 'NBANDSGW', 'NBANDSLF', 'DIM', 'ANTIRES', 'LUSEW', 'OMEGAGRID', 'SELFENERGY', 'NKREDLFX', 'NKREDLFY',
'NKREDLFZ', 'MAXMEM', 'TELESCOPE', 'LCRITICAL_MEM', 'GGA2',
'TURBO', 'QUAD_EFG','IRESTART','NREBOOT','NMIN','EREF','KSPACING','KGAMMA','LSUBROT','SCALEE','LVHAR','LORBITALREAL','DARWINR','DARWINV','LFOCKAEDFT','NUCIND','MAGPOS','LNICSALL','LADDER','LHARTREE','IBSE','NBANDSO','NBANDSV','OPTEMAX')

class Incar(dict, VaspInput):
    """
    INCAR object for reading and writing INCAR files
    essentially consists of a dictionary with some helper functions
    """
    
    def __init__(self,params = dict()):
        """
        Creates an Incar object.
        Optional arguments:
            params - A set of input parameters as a dictionary.
        """
        super(Incar,self).__init__()
        self.update(params)

    def __setitem__(self, key, val):
        """
        Add parameter-val pair to Incar.  Warns if parameter is not in list of valid INCAR tags.
        Also cleans the parameter and val by stripping leading and trailing white spaces.
        """
        if key.strip().upper() not in VALID_INCAR_TAGS:
            warnings.warn(key.strip() + " not in VALID_INCAR_TAGS")
        super(Incar,self).__setitem__(key.strip(), Incar.proc_val(key.strip(), val.strip()) if isinstance(val,basestring) else val)

    def get_string(self, sort_keys = False, pretty = False):
        keys = self.keys()
        if sort_keys:
            keys = sorted(keys)
        lines = []
        for k in keys:
            if isinstance(self[k], list):
                lines.append([k," ".join([str(i) for i in self[k]])])
            else:
                lines.append([k,self[k]])
        
        if pretty:
            return str_aligned(lines)
        else:
            return str_delimited(lines, None," = ")
    
    @staticmethod
    def from_structure(structure, parameter_set = VaspParameterSet()):
        incar = Incar()
        poscar = Poscar(structure)
        for key, setting in parameter_set.incar_settings.items():
            if key == "MAGMOM":
                incar[key] =  [setting.get(site.specie.symbol, 0.6) for site in structure]
            elif key in ['LDAUU', 'LDAUJ', 'LDAUL']:
                incar[key] =  [setting.get(sym, 0) for sym in poscar.site_symbols]
            else:
                incar[key] = setting
                
        has_u = sum(incar['LDAUU']) > 0
        if not has_u:
            for key in incar.keys():
                if key.startswith('LDAU'):
                    del incar[key]
        
        return incar
    
    def __str__(self):
        return self.get_string(sort_keys = True, pretty = False)

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename, "r") as f:
            lines = list(clean_lines(f.readlines()))
        params = {}
        for line in lines:
            m = re.match("(\w+)\s*=\s*(.*)", line)
            if m:
                key = m.group(1).strip()
                val = m.group(2).strip()
                val = Incar.proc_val(key, val)
                params[key] = val
        return Incar(params)
    
    @staticmethod
    def proc_val(key, val):
        list_type_keys = ('LDAUU', 'LDAUL', 'LDAUJ', 'LDAUTYPE','MAGMOM')
        boolean_type_keys = ('LDAU', 'LWAVE')
        number_type_keys = ('NSW', 'NELMIN', 'ISIF', 'IBRION', "ISPIN", "EDIFF", "ICHARG", "NELM", "ISMEAR", "NPAR", "SIGMA", "LDAUPRINT", 'LMAXMIX')
         
        
        def smart_int_or_float(numstr):
            if numstr.find(".") != -1:
                return float(numstr)
            else:
                return int(numstr)
        try:
            if key in list_type_keys:
                output = list()
                toks = re.split("\s+", val)
                
                for tok in toks:
                    m = re.match("(\d+)\*([\d\.\-\+]+)", tok)
                    if m:
                        output.extend([smart_int_or_float(m.group(2))] * int(m.group(1)))
                    else:
                        output.append(smart_int_or_float(tok))
                return output
            if key in boolean_type_keys:
                m = re.search("^\W+([TtFf])", val)
                if m:
                    if m.group(1) == "T" or m.group(1) == "t":
                        return True
                    else:
                        return False
                raise ValueError(key + " should be a boolean type!")
            
            if key in number_type_keys:
                return smart_int_or_float(val)
        except:
            return val
              
        return val
            
    def diff(self, other):
        similar_param = {}
        different_param = {}
        for k1,v1 in self.items():
            if k1 not in other:
                different_param[k1] = {"INCAR1": v1, "INCAR2": 'Default'}
            elif v1 != other[k1]:
                different_param[k1] = {"INCAR1": v1, "INCAR2": other[k1]}
            else:
                similar_param[k1] = v1
        for k2,v2 in other.items():
            if k2 not in similar_param and k2 not in different_param:
                if k1 not in self:
                    different_param[k1] = {"INCAR1": 'Default', "INCAR2": v2}
        return {'Same parameters' : similar_param, 'Different': different_param}
                
    def __add__(self, other):
        """
        Add all the values of another INCAR object to this object
        Facilitates the use of "standard" INCARs
        """
        params = {k:v for k,v in self.items()}
        for k, v in other.items():
            if k in self and v != self[k]:
                raise ValueError("Incars have conflicting values!")
            else:
                params[k] = v
        return Incar(params)

class Kpoints(VaspInput):
    """
    Very basic KPOINT reader/writer
    The lines of the original file
    are stored as is but the kpts are stored
    as a list which can be accessed and
    modified
    GH: NEEDS to be updated to work with band structure kpoint files
    """
    def __init__(self, comment = "Default", num_kpts = 0, style = "Gamma", kpts = [[1,1,1]], kpts_shift = [0,0,0]):
        self.comment = comment
        self.num_kpts = num_kpts
        self.style = style
        self.kpts = kpts
        self.kpts_shift = kpts_shift

    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename) as f:
            lines = list(clean_lines(f.readlines(), False))
        kpoints = Kpoints()
        kpoints.comment = lines[0].strip()
        kpoints.num_kpts = int(lines[1].strip())
        style = lines[2].strip().lower()[0]
        if style == "a":
            kpoints.style = "Auto"
            kpoints.kpts = [[int(lines[3].strip())]]
        elif style == "g" or style == "m":
            kpoints.style = "Gamma" if style == "g" else "Monkhorst-Pack"
            kpoints.kpts = [[int(x) for x in lines[3].strip().split()]]
            if len(lines)>4 and coord_pattern.match(lines[4].strip()):
                try:
                    kpoints.kpts_shift = [int(x) for x in lines[4].strip().split()]
                except:
                    pass
        elif kpoints.num_kpts <= 0:
            kpoints.style = "Cartesian" if style == "c" else "Reciprocal"
            kpoints.kpts = [[float(x) for x in lines[i].strip().split()] for i in xrange(3,6)]
            kpoints.kpts_shift = [float(x) for x in lines[6].strip().split()]
                
            
        return kpoints

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    def __str__(self):
        lines = []
        lines += [self.comment]
        lines += [str(self.num_kpts)]
        lines += [self.style]
        for line in self.kpts:
            lines += [" ".join([str(x) for x in line])]
        lines += [" ".join([str(x) for x in self.kpts_shift])]
        return "\n".join(lines)
    
    @property
    def to_dict(self):
        d = {'comment': self.comment, 'nkpoints' : self.num_kpts, 'generation_style' : self.style, 'kpoints': self.kpts, 'usershift': self.kpts_shift}
        optional_paras = ['genvec1', 'genvec2', 'genvec3', 'shift']
        for para in optional_paras:
            if para in self.__dict__:
                d[para] = self.__dict__[para]
        return d 
    
class PotcarSingle(VaspInput):
    """
    Object for a **single** POTCAR.
    The builder assumes the complete string is the POTCAR
    contains the complete untouched data in "data" as a string and
    a dictionary of keywords
    """
    def __init__(self, data):
        """
        Expects a complete and single potcar file as a string in "data"
        """
        self.data = data # raw POTCAR as a string
        keyValPairs = re.compile(r";*\s*(.+?)\s*=\s*([^;\n]+)\s*", re.M).findall(data)
        self.keywords = dict(keyValPairs) # all key = val found in the POTCAR as a dictionary all keys and vals are strings

    def __str__(self):
        return self.data

    def write_file(self, filename):
        writer = open(filename, 'w')
        writer.write(self.__str__() + "\n")
        writer.close()

    @property
    def symbol(self):
        """
        Full name of POTCAR, e.g., Fe_pv
        """
        return self.keywords['TITEL'].split(" ")[1].strip()        
    
    @property
    def element(self):
        """
        Attempt to return the atomic symbol based on the VRHFIN keyword
        """
        return self.keywords['VRHFIN'].split(":")[0].strip()

    @property
    def atomic_no(self):
        """
        Attempt to return the atomic number based on the VRHFIN keyword
        """
        return Element(self.element).Z

class Potcar(list,VaspInput):
    """
    Object for reading and writing POTCAR files for
    calculations.
    """
    functional_dir = {'PBE':'POT_GGA_PAW_PBE', 'LDA':'POT_LDA_PAW', 'PW91':'POT_GGA_PAW_PW91'}
    
    def __init__(self, symbols = None):
        if symbols != None:
            self.set_symbols(symbols)

    @staticmethod
    def from_file(filename):
        with file_open_zip_aware(filename, "r") as reader:
            fData = reader.read()
        potcar = Potcar()
        potcar_strings = re.compile(r"\n{0,1}\s*(.*?End of Dataset)", re.S).findall(fData)
        for p in potcar_strings:
            potcar.append(PotcarSingle(p))
        return potcar

    def __str__(self):
        return "\n".join([str(potcar) for potcar in self])

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(self.__str__() + "\n")
        
    @property
    def symbols(self):
        """
        Get the atomic symbols of all the atoms in the POTCAR file
        """
        return [p.symbol for p in self]

    def set_symbols(self, elements, functional = 'PBE', parameter_set = VaspParameterSet()):
        module_dir = os.path.dirname(pymatgen.__file__)
        config = ConfigParser.SafeConfigParser()
        config.readfp(open(os.path.join(module_dir, "pymatgen.cfg")))
        VASP_PSP_DIR = os.path.join(config.get('VASP', 'pspdir'), Potcar.functional_dir[functional])
        
        del self[:]
        for el in elements:
            sym = el
            if parameter_set and Element.is_valid_symbol(el) and el in parameter_set.potcar_settings:
                sym = parameter_set.potcar_settings[el]
            reader = file_open_zip_aware(os.path.join(VASP_PSP_DIR, "POTCAR." + sym + ".gz"))
            self.append(PotcarSingle(reader.read()))
            reader.close()

class Vasprun(object):
    """
    Vastly improved sax-based parser for vasprun.xml files.
    Speedup over Dom is at least 2x for smallish files (~1Mb) to orders of magnitude for larger files (~10Mb).
    All data is stored as attributes, which are delegated to the VasprunHandler object.
    Accessible attributes from VasprunHandler are:
        Vasp results
        ------------
        energies - All energies in run, represented as a list of ionic steps, with a list of scstep energies.
                   E.g. [ [0.1, 0.2, 0.1], [0.1], [0.22, 0.23, 0.20, 0.205]] 
        structures - List of Structure objects for the structure at each ionic step.
        tdos - Total dos calculated at the end of run.
        idos - Integrated dos calculated at the end of run.
        pdos - List of list of PDos objects. Access as pdos[atomindex][orbitalindex]
        efermi - Fermi energy
        eigenvalues - Final eigenvalues as a dict of {(kpoint index, Spin.up):[[eigenvalue, occu]]}. 
                      This representation is probably not ideal, but since this is not used anywhere else for now, I leave it as such.
                      Future developers who need to work with this should refactored the object into a sensible structure.
        
        Vasp inputs
        -----------
        incar - Incar object for parameters specified in INCAR file.
        parameters - Incar object with parameters that vasp actually used, including all defaults.
        kpoints = Kpoints object for KPOINTS specified in run.
        actual_kpoints - List of actual kpoints, e.g., [[0.25, 0.125, 0.08333333], [-0.25, 0.125, 0.08333333], [0.25, 0.375, 0.08333333], ....]
        actual_kpoints_weights = List of kpoint weights, E.g., [0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, 0.04166667, ....]
        atomic_symbols - List of atomic symbols, e.g., [u'Li', u'Fe', u'Fe', u'P', u'P', u'P']
        potcar_symbols - List of POTCAR symbols. E.g., [u'PAW_PBE Li 17Jan2003', u'PAW_PBE Fe 06Sep2000', ..]
    
    A few helper attributes have also been added to get commonly used results such as final energies.
    
    Author: <NAME>
    """
    supported_properties = ['lattice_rec','vasp_version', 'incar', 'parameters', 'potcar_symbols', 'atomic_symbols', 'kpoints', 'actual_kpoints', 'structures',
                            'actual_kpoints_weights', 'dos_energies', 'eigenvalues', 'tdos', 'idos', 'pdos', 'efermi', 'ionic_steps', 'dos_has_errors']
    
    def __init__(self, filename):
        self._filename = filename    
        with file_open_zip_aware(filename) as f:
            self._handler = VasprunHandler(filename)
            self._parser = xml.sax.parse(f, self._handler)
            for k in Vasprun.supported_properties:
                setattr(self, k, getattr(self._handler, k))
    
    @property
    def converged(self):
        return len(self.structures) - 2 < self.parameters['NSW'] or self.parameters['NSW'] == 0
    
    @property
    def final_energy(self):
        """
        Final energy from the vasp run.
        """
        return self.ionic_steps[-1]['electronic_steps'][-1]['e_wo_entrp']
    
    @property
    def final_structure(self):
        """
        Final structure from vasprun.
        """
        return self.structures[-1]
    
    @property
    def initial_structure(self):
        """
        Initial structure from vasprun.
        """
        return self.structures[0]
    
    @property
    def complete_dos(self):
        """
        A complete dos object which incorporates the total dos and all projected dos.
        """
        return CompleteDos(self.final_structure, self.tdos, self.pdos)
    
    @property
    def to_dict(self):
        d = {}
        d['vasp_version'] = self.vasp_version
        d['has_vasp_completed'] = self.converged
        d['nsites'] = len(self.final_structure)
        d['unit_cell_formula'] = self.final_structure.composition.to_dict
        comp = self.final_structure.composition
        d['reduced_cell_formula'] = Composition.from_formula(comp.reduced_formula).to_dict
        d['pretty_formula'] = comp.reduced_formula
        symbols = [re.split("\s+", s)[1] for s in self.potcar_symbols]
        symbols = [re.split("_", s)[0] for s in symbols]
        d['elements'] = symbols
        d['nelements'] = len(symbols)
        d['is_hubbard'] = self.incar.get('LDAU', False)
        if d['is_hubbard']:
            us = self.incar.get('LDAUU', self.parameters.get('LDAUU'))
            js = self.incar.get('LDAUJ', self.parameters.get('LDAUJ'))
            if len(us) == len(symbols):
                d['hubbards'] = { symbols[i] : us[i] - js[i] for i in xrange(len(symbols))}
            elif sum(us) == 0 and sum(js) == 0:
                d['is_hubbard'] = False
                d['hubbards'] = {}
            else:
                raise VaspParserError("Length of U value parameters and atomic symbols are mismatched")
        else:
            d['hubbards'] = {}
        if d['is_hubbard']:
            d['run_type'] = "GGA+U"
        elif self.parameters.get('LHFCALC', False):
            d['run_type'] = "HF"
        else:
            d['run_type'] = "GGA"
                   
        d['input'] = {}
        d['input']['incar'] = {k:v for k,v in self.incar.items()}
        d['input']['crystal'] = self.initial_structure.to_dict
        d['input']['kpoints'] = self.kpoints.to_dict
        d['input']['kpoints']['actual_points'] = [{'abc':list(self.actual_kpoints[i]), 'weight':self.actual_kpoints_weights[i]} for i in xrange(len(self.actual_kpoints))] 
        d['input']['potcar'] = [s.split(" ")[1] for s in self.potcar_symbols]
        d['input']['parameters'] = {k:v for k,v in self.parameters.items()}
        d['input']['lattice_rec'] = self.lattice_rec
        
        d['output'] = {}
        d['output']['ionic_steps'] = clean_json(self.ionic_steps)
        d['output']['final_energy'] = self.final_energy
        d['output']['final_energy_per_atom'] = self.final_energy / len(self.final_structure)
        d['output']['crystal'] = self.final_structure.to_dict
        d['output']['efermi'] = self.efermi
        #{(kpoint index, Spin.up):array(float)}
        
        d['output']['eigenvalues'] = {}
        for (index, spin), values in self.eigenvalues.items():
            if str(index) not in d['output']['eigenvalues']:
                d['output']['eigenvalues'][str(index)] = {str(spin):values}
            else:
                d['output']['eigenvalues'][str(index)][str(spin)] = values
                
        return d

 
class VasprunHandler(xml.sax.handler.ContentHandler):
    """
    Sax handler for vasprun.xml.
    Attributes are mirrored into Vasprun object.
    
    Author: <NAME>
    """
        
    def __init__(self, filename):
        self.filename = filename
        # variables to be filled
        self.vasp_version = None
        self.incar = Incar()
        self.parameters = Incar()
        self.potcar_symbols = []
        self.atomic_symbols = []
        self.kpoints = Kpoints()
        self.actual_kpoints = []
        self.actual_kpoints_weights = []
        self.dos_energies = None
        self.eigenvalues = {}#  will  be  {(kpoint index, Spin.up):array(float)}
        self.tdos = {}
        self.idos = {}
        self.pdos = {}
        self.efermi = None 
        self.ionic_steps = [] # should be a list of dict
        self.structures = []
        self.lattice_rec = []
        
        self.input_read = False
        self.all_calculations_read = False
        self.read_structure = False
        self.read_rec_lattice = False
        self.read_calculation = False
        self.read_eigen = False
        self.read_dos= False
        self.in_efermi = False
        self.read_atoms = False
        self.read_lattice = False
        self.read_positions = False
        self.incar_param = None
        self.dos_energies_val = []
        self.dos_val = []
        self.idos_val = []
        self.raw_data = []
        self.dos_has_errors = False #will be set to true if there is an error parsing the Dos.
        self.state = defaultdict(bool)
        
    
    def in_all(self, xml_tags):
        return all([getattr(self, 'in_'+tag, None) for tag in xml_tags])
    
    def startElement(self, name, attributes):
        
        self.state[name] = True if 'name' not in attributes else attributes['name']
        self.read_val = False
        
        #Nested if loops makes reading much faster.
        if not self.input_read: #reading input parameters
            if (name == "i" or name == "v") and (self.state['incar'] or self.state['parameters']):
                self.incar_param = attributes['name']
                self.param_type = 'float' if 'type' not in attributes else attributes['type']
                self.read_val = True
            elif name == "v" and self.state['kpoints']:
                self.read_val = True
            elif name == "generation" and self.state['kpoints']:
                self.kpoints.comment   = "Kpoints from vasprun.xml"
                self.kpoints.num_kpts  = 0
                self.kpoints.style = attributes['param']
                self.kpoints.kpts = []
                self.kpoints.kpts_shift = [0,0,0]
            elif name == "c" and (self.state['array'] == "atoms" or self.state['array'] == "atomtypes"):
                self.read_val = True
            elif name == "i" and self.state['i'] == "version" and self.state['generator']:
                self.read_val = True
                
        else: #reading calculations and structures.
            if self.read_calculation:
                if name == "i" and self.state['scstep']:
                    self.read_val = True
                elif name == "v" and (self.state['varray'] == "forces" or self.state['varray'] == "stress"):
                    self.read_positions = True
            if self.read_structure:
                if name == "v" and self.state['varray'] == 'basis':
                    self.read_lattice = True
                elif name == "v" and self.state['varray'] == 'positions':
                    self.read_positions = True
                if name == 'v' and self.state['varray'] == 'rec_basis':
                    self.read_rec_lattice = True
            if name == "calculation":
                self.scdata = []
                self.read_calculation = True
            elif name == "scstep":
                self.scstep = {}
            elif name == 'structure':
                self.latticestr = StringIO.StringIO()
                self.latticerec = StringIO.StringIO()
                self.posstr = StringIO.StringIO()
                self.read_structure = True
            elif name == 'varray' and (self.state['varray'] == "forces" or self.state['varray'] == "stress"):
                self.posstr = StringIO.StringIO()
            elif name == "eigenvalues":
                self.all_calculations_read = True
            if self.read_eigen:
                if name == "r" and self.state["set"]:
                    self.read_val = True
                elif name == "set" and "comment" in attributes:
                    comment = attributes["comment"]
                    self.state["set"] = comment
                    if comment.startswith("spin"):
                        self.eigen_spin = Spin.up if self.state["set"] == "spin 1" else Spin.down
                    if comment.startswith("kpoint"):
                        self.eigen_kpoint = int(comment.split(" ")[1])
            elif self.read_dos:
                if (name == "i" and self.state["i"] == "efermi") or (name == "r" and self.state["set"] ):
                    self.read_val = True
                elif name == "set" and "comment" in attributes:
                    comment = attributes["comment"]
                    self.state["set"] = comment
                    if self.state['partial']:
                        if comment.startswith("ion"):
                            self.pdos_ion = int(comment.split(" ")[1])
                        elif comment.startswith("spin"):
                            self.pdos_spin = Spin.up if self.state["set"] == "spin 1" else Spin.down
            elif name == "dos":
                self.dos_energies = None
                self.tdos = {}
                self.idos = {}
                self.pdos = {}
                self.efermi = None 
                self.read_dos = True
            elif name == "eigenvalues":
                self.eigenvalues = {}#  will  be  {(kpoint index, Spin.up):array(float)}
                self.read_eigen = True
            
        if self.read_val:
            self.val = StringIO.StringIO()

    def characters(self, data):
        if self.read_val:
            self.val.write(data)
        if self.read_lattice:
            self.latticestr.write(data)
        elif self.read_positions:
            self.posstr.write(data)
        elif self.read_rec_lattice:
            self.latticerec.write(data)

    
    #To correct for stupid vasp bug which names Xenon as X!!
    EL_MAPPINGS = {'X':'Xe'}
    
    def endElement(self, name):
        
        if not self.input_read:
            if name == "i":
                if self.state['incar']:
                    self.incar[self.incar_param] = parse_parameters(self.param_type, self.val.getvalue().strip())
                elif self.state['parameters']:
                    self.parameters[self.incar_param] = parse_parameters(self.param_type, self.val.getvalue().strip())
                elif self.state['generator'] and self.state["i"] == "version":
                    self.vasp_version = self.val.getvalue().strip()
                self.incar_param = None
            elif name == "set":
                if self.state['array'] == "atoms":
                    self.atomic_symbols = self.atomic_symbols[::2]
                    self.atomic_symbols = [sym if sym not in VasprunHandler.EL_MAPPINGS else VasprunHandler.EL_MAPPINGS[sym] for sym in self.atomic_symbols]
                elif self.state['array'] == "atomtypes":
                    self.potcar_symbols = self.potcar_symbols[4::5]
                    self.input_read = True
            elif name == "c":
                if self.state['array'] == "atoms":
                    self.atomic_symbols.append(self.val.getvalue().strip())
                elif self.state['array'] == "atomtypes":
                    self.potcar_symbols.append(self.val.getvalue().strip())
            elif name == "v":
                if self.state['incar']:
                    self.incar[self.incar_param] = parse_v_parameters(self.param_type, self.val.getvalue().strip(), self.filename, self.incar_param)
                    self.incar_param = None
                elif self.state['parameters']:
                    self.parameters[self.incar_param] = parse_v_parameters(self.param_type, self.val.getvalue().strip(), self.filename, self.incar_param)
                elif self.state['kpoints']:
                    if self.state['varray'] == 'kpointlist':
                        self.actual_kpoints.append([float(x) for x in re.split("\s+",self.val.getvalue().strip())])
                    if self.state['varray'] == 'weights':
                        self.actual_kpoints_weights.append(float(self.val.getvalue()))
                    if self.state['v'] == "divisions":
                        self.kpoints.kpts   = [[int(x) for x in re.split("\s+",self.val.getvalue().strip())]]
                    elif self.state['v'] == "usershift":
                        self.kpoints.kpts_shift   = [float(x) for x in re.split("\s+",self.val.getvalue().strip())]
                    elif self.state['v'] == "genvec1" or self.state['v'] == "genvec2" or self.state['v'] == "genvec3" or self.state['v'] == "shift":
                        setattr(self.kpoints, self.state['v'], [float(x) for x in re.split("\s+",self.val.getvalue().strip())])
                        
        else:
            if self.read_calculation:
                if name == "i" and self.state['scstep']:
                    self.scstep[self.state['i']] = float(self.val.getvalue())
                elif name == 'scstep':
                    self.scdata.append(self.scstep)
                elif name == 'varray' and self.state['varray'] == "forces":
                    self.forces = np.array([float(x) for x in re.split("\s+",self.posstr.getvalue().strip())])
                    self.forces.shape = (len(self.atomic_symbols), 3)
                elif name == 'varray' and self.state['varray'] == "stress":
                    self.stress = np.array([float(x) for x in re.split("\s+",self.posstr.getvalue().strip())])
                    self.stress.shape = (3, 3)
                elif name == "calculation":
                    self.ionic_steps.append({'electronic_steps':self.scdata, 'structure':self.structures[-1], 'forces': self.forces, 'stress':self.stress})
                    self.read_calculation = False
            if self.read_structure:
                if name == "v":
                    self.read_positions = False
                    self.read_lattice = False
                    self.read_lattice_rec = False
                    self.read_rec_lattice = False
                elif name == "structure":
                    self.lattice = np.array([float(x) for x in re.split("\s+",self.latticestr.getvalue().strip())])
                    self.lattice.shape = (3,3)
                    self.pos = np.array([float(x) for x in re.split("\s+",self.posstr.getvalue().strip())])
                    self.pos.shape = (len(self.atomic_symbols), 3)
                    self.structures.append(Structure(self.lattice, self.atomic_symbols, self.pos))
                    self.lattice_rec = np.array([float(x) for x in re.split("\s+",self.latticerec.getvalue().strip())])
                    self.lattice_rec.shape = (3,3)
                    self.read_structure = False
            elif self.read_dos:
                try:
                    if name == "i" and self.state["i"] == "efermi":
                        self.efermi = float(self.val.getvalue().strip())
                    elif name == "r" and self.state["total"]  and str(self.state["set"]).startswith("spin"):
                        tok = re.split("\s+", self.val.getvalue().strip())
                        self.dos_energies_val.append(float(tok[0]))
                        self.dos_val.append(float(tok[1]))
                        self.idos_val.append(float(tok[2]))
                    elif name == "r" and self.state["partial"]  and str(self.state["set"]).startswith("spin"):
                        tok = re.split("\s+", self.val.getvalue().strip())
                        self.raw_data.append([float(i) for i in tok[1:]])
                    elif name == "set" and self.state["total"] and str(self.state["set"]).startswith("spin"):
                        spin = Spin.up if self.state["set"] == "spin 1" else Spin.down
                        self.tdos[spin] = self.dos_val
                        self.idos[spin] = self.dos_val
                        self.dos_energies = self.dos_energies_val
                        self.dos_energies_val = []
                        self.dos_val = []
                        self.idos_val = []
                    elif name == "set" and self.state["partial"] and str(self.state["set"]).startswith("spin"):
                        spin = Spin.up if self.state["set"] == "spin 1" else Spin.down
{'fraction_non_alphanumeric': 0.13326813138193208, 'fraction_numerical': 0.007381792355071198, 'mean_word_length': 1.9692318329530527, 'pattern_counts': {'":': 68, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 17}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4972409', 'n_tokens_mistral': 27294, 'n_tokens_neox': 26008, 'n_words': 6480}