# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.

import os
import re

from ase.atoms import Atoms
from ase.constraints import FixCartesian
import numpy as np
from collections import OrderedDict
from pyiron_vasp.dft.waves.electronic import ElectronicStructure
import defusedxml.cElementTree as ETree
from defusedxml.ElementTree import ParseError
import warnings


__author__ = "Sudarsan Surendralal"
__copyright__ = (
    "Copyright 2021, Max-Planck-Institut für Eisenforschung GmbH - "
    "Computational Materials Design (CM) Department"
)
__version__ = "1.0"
__maintainer__ = "Sudarsan Surendralal"
__email__ = "surendralal@mpie.de"
__status__ = "production"
__date__ = "Sep 1, 2017"


class Vasprun(object):
    """
    This module is used to parse vasprun.xml files and store the data consistent with the pyiron input/output storage
    formats.

    Attributes:

        vasprun_dict (dict): Dictionary containing all information from the calculation parsed from the vasprun.xml
                            file. If you consider a simulation with N atoms and M ionic steps

                'positions' (numpy.ndarray): MxNx3 array containing all the relative positions
                'cell' (numpy.ndarray): Mx3x3 array containing all the size and shape of cells at every iteration point
                'forces' (numpy.ndarray): MxNx3 array containing all the forces in eV/A
                'total_energies' (numpy.ndarray): 1xM array containing all the total energies in eV
    """

    def __init__(self):
        self.vasprun_dict = dict()

    def from_file(self, filename="vasprun.xml"):
        """
        Parsing vasprun.xml from the working directory

        Args:
            filename (str): Path to the vasprun file
        """
        if not (os.path.isfile(filename)):
            raise AssertionError()
        try:
            self.parse_root_to_dict(filename)
        except ParseError:
            raise VasprunError(
                "The vasprun.xml file is either corrupted or the simulation has failed"
            )

    def parse_root_to_dict(self, filename):
        """
        Parses from the main xml root.
        """
        d = self.vasprun_dict
        d["scf_energies"] = list()
        d["scf_fr_energies"] = list()
        d["scf_0_energies"] = list()
        d["scf_dipole_moments"] = list()
        d["positions"] = list()
        d["cells"] = list()
        d["forces"] = list()
        d["total_energies"] = list()
        d["total_fr_energies"] = list()
        d["total_0_energies"] = list()
        d["kinetic_energies"] = list()
        d["stress_tensors"] = list()
        for _, leaf in ETree.iterparse(filename):
            if leaf.tag in ["generator", "incar"]:
                d[leaf.tag] = dict()
                for items in leaf:
                    d[leaf.tag] = self.parse_item_to_dict(items, d[leaf.tag])
            if leaf.tag in ["kpoints"]:
                d[leaf.tag] = dict()
                self.parse_kpoints_to_dict(leaf, d[leaf.tag])
            if leaf.tag in ["atominfo"]:
                d[leaf.tag] = dict()
                self.parse_atom_information_to_dict(leaf, d[leaf.tag])
            if leaf.tag in ["structure"] and "name" in leaf.keys():
                if "initialpos" in leaf.attrib["name"]:
                    d["init_structure"] = dict()
                    self.parse_structure_to_dict(leaf, d["init_structure"])
                elif "finalpos" in leaf.attrib["name"]:
                    d["final_structure"] = dict()
                    self.parse_structure_to_dict(leaf, d["final_structure"])
            if leaf.tag in ["calculation"]:
                self.parse_calc_to_dict(leaf, d)
            if leaf.tag in ["parameters"]:
                pass
                self.parse_parameters(leaf, d)
        d["cells"] = np.array(d["cells"])
        d["positions"] = np.array(d["positions"])
        # Check if the parsed coordinates are in absolute/relative coordinates. If absolute, convert to relative
        total_positions = d["positions"].flatten()
        if len(np.argwhere(total_positions > 1)) / len(total_positions) > 0.2:
            pos_new = d["positions"].copy()
            for i, pos in enumerate(pos_new):
                d["positions"][i] = np.dot(pos, np.linalg.inv(d["cells"][i]))
        d["forces"] = np.array(d["forces"])
        d["total_energies"] = np.array(d["total_energies"])
        d["total_fr_energies"] = np.array(d["total_fr_energies"])
        d["total_0_energies"] = np.array(d["total_0_energies"])
        if len(d["kinetic_energies"]) > 0:
            d["kinetic_energies"] = np.array(d["kinetic_energies"])
        else:
            del d["kinetic_energies"]
        d["scf_energies"] = d["scf_energies"]
        d["scf_dipole_moments"] = d["scf_dipole_moments"]
        d["scf_fr_energies"] = d["scf_fr_energies"]
        d["scf_0_energies"] = d["scf_0_energies"]
        d["stress_tensors"] = d["stress_tensors"]

    def parse_kpoints_to_dict(self, node, d):
        """
        Parses k-points data from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "kpoints"):
            raise AssertionError()
        for leaf in node:
            if leaf.tag == "generation":
                d[leaf.tag] = dict()
                d[leaf.tag]["scheme"] = leaf.attrib["param"]
                if d[leaf.tag]["scheme"] == "listgenerated":
                    line_mode_kpoints = list()
                    for item in leaf:
                        if item.tag == "v":
                            line_mode_kpoints.append(
                                self._parse_vector(item, vec_type=float)
                            )
                    d["line_mode_kpoints"] = np.array(line_mode_kpoints)
                else:
                    gen_vec = np.zeros((3, 3))
                    for item in leaf:
                        if item.tag == "v":
                            if item.attrib["name"] in ["divisions"]:
                                d[leaf.tag]["divisions"] = self._parse_vector(
                                    item, vec_type=int
                                )
                            if item.attrib["name"] in [
                                "genvec{}".format(i) for i in range(1, 4)
                            ]:
                                if item.attrib["name"] == "genvec1":
                                    gen_vec[0, :] = self._parse_vector(
                                        item, vec_type=float
                                    )
                                if item.attrib["name"] == "genvec2":
                                    gen_vec[1, :] = self._parse_vector(
                                        item, vec_type=float
                                    )
                                if item.attrib["name"] == "genvec3":
                                    gen_vec[2, :] = self._parse_vector(
                                        item, vec_type=float
                                    )
                            if item.attrib["name"] in ["shift", "usershift"]:
                                d[leaf.tag][item.attrib["name"]] = self._parse_vector(
                                    item, vec_type=float
                                )
                    d[leaf.tag]["genvec"] = np.array(gen_vec)
            if leaf.tag == "varray":
                if leaf.attrib["name"] == "kpointlist":
                    d["kpoint_list"] = self._parse_2d_matrix(leaf, vec_type=float)
                if leaf.attrib["name"] == "weights":
                    d["kpoint_weights"] = self._parse_2d_matrix(
                        leaf, vec_type=float
                    ).flatten()

    def parse_atom_information_to_dict(self, node, d):
        """
        Parses atom information from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "atominfo"):
            raise AssertionError()
        species_dict = OrderedDict()
        for leaf in node:
            if leaf.tag == "atoms":
                d["n_atoms"] = self._parse_vector(leaf)[0]
            if leaf.tag == "types":
                d["n_species"] = self._parse_vector(leaf)[0]
            if leaf.tag == "array":
                if leaf.attrib["name"] == "atomtypes":
                    for item in leaf:
                        if item.tag == "set":
                            for sp in item:
                                elements = sp
                                if elements[1].text in species_dict.keys():
                                    count = 1
                                    not_unique = True
                                    species_key = None
                                    while not_unique:
                                        species_key = "_".join(
                                            [elements[1].text, str(count)]
                                        )
                                        if species_key not in species_dict.keys():
                                            not_unique = False
                                        else:
                                            count += 1
                                    if species_key is not None:
                                        species_dict[species_key] = dict()
                                        species_dict[species_key]["n_atoms"] = int(
                                            elements[0].text
                                        )
                                        species_dict[species_key]["valence"] = float(
                                            elements[3].text
                                        )
                                else:
                                    species_key = elements[1].text
                                    species_dict[species_key] = dict()
                                    species_dict[species_key]["n_atoms"] = int(
                                        elements[0].text
                                    )
                                    species_dict[species_key]["valence"] = float(
                                        elements[3].text
                                    )
        d["species_dict"] = species_dict
        species_list = list()
        for key, val in species_dict.items():
            for sp in np.tile([key], species_dict[key]["n_atoms"]):
                species_list.append(clean_character(sp))
        d["species_list"] = species_list

    def parse_fermi_level_to_dict(self, node, d):
        """
        Parses fermi level from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "dos"):
            raise AssertionError()
        for item in node:
            if item.tag == "i":
                self.parse_item_to_dict(item, d)

    def parse_total_dos_to_dict(self, node, d):
        """
        Parses total dos data from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "total"):
            raise AssertionError()
        for item in node:
            if item.tag == "array":
                for ii in item:
                    if ii.tag == "set":
                        spin_dos_energies = list()
                        spin_dos_density = list()
                        spin_dos_idensity = list()
                        for sp in ii:
                            if sp.tag == "set" and "spin" in sp.attrib["comment"]:
                                try:
                                    values = self._parse_2d_matrix(sp, vec_type=float)
                                    dos_energies = values[:, 0]
                                    dos_density = values[:, 1]
                                    dos_idensity = values[:, 2]
                                    spin_dos_energies.append(dos_energies)
                                    spin_dos_density.append(dos_density)
                                    spin_dos_idensity.append(dos_idensity)
                                except ValueError:
                                    pass
                        d["spin_dos_energies"] = np.array(spin_dos_energies)
                        d["spin_dos_density"] = np.array(spin_dos_density)
                        d["spin_dos_idensity"] = np.array(spin_dos_idensity)

    def parse_partial_dos_to_dict(self, node, d):
        """
        Parses partial dos data from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "partial"):
            raise AssertionError()
        orbital_dict = dict()
        orbital_index = 0
        for item in node:
            if item.tag == "array":
                for ii in item:
                    if ii.tag == "field":
                        if "energy" not in ii.text:
                            orbital_dict[ii.text.replace(" ", "")] = orbital_index
                        orbital_index += 1
                    if ii.tag == "set":
                        atom_resolved_dos = list()
                        for ion in ii:
                            spin_resolved_dos = list()
                            if ion.tag == "set" and "ion" in ion.attrib["comment"]:
                                for sp in ion:
                                    if (
                                        sp.tag == "set"
                                        and "spin" in sp.attrib["comment"]
                                    ):
                                        values = self._parse_2d_matrix(
                                            sp, vec_type=float
                                        )
                                        spin_resolved_dos.append(values[:, 1:])
                            atom_resolved_dos.append(spin_resolved_dos)
                        atom_resolved_dos = np.array(atom_resolved_dos)
                        n_atoms, n_spin, n_densities, n_orbitals = np.shape(
                            atom_resolved_dos
                        )
                        new_grand_dos_matrix = np.zeros(
                            (n_spin, n_atoms, n_orbitals, n_densities)
                        )
                        for i_spin in range(n_spin):
                            for i_atom in range(n_atoms):
                                for i_orbitals in range(n_orbitals):
                                    new_grand_dos_matrix[
                                        i_spin, i_atom, i_orbitals, :
                                    ] = atom_resolved_dos[i_atom, i_spin, :, i_orbitals]
                        d["resolved_dos_matrix"] = new_grand_dos_matrix

    def parse_projected_dos_to_dict(self, node, d):
        """
        Parses partial dos data from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "projected"):
            raise AssertionError()
        orbital_dict = dict()
        orbital_index = 0
        for item in node:
            if item.tag == "array":
                for ii in item:
                    if ii.tag == "field":
                        orbital_dict[ii.text] = orbital_index
                        orbital_index += 1
                    if ii.tag == "set":
                        spin_dos_mat = list()
                        for sp in ii:
                            if sp.tag == "set" and "spin" in sp.attrib["comment"]:
                                kpt_dos_mat = list()
                                for kpt in sp:
                                    band_dos_mat = list()
                                    for band in kpt:
                                        dos_matrix = self._parse_2d_matrix(
                                            band, vec_type=float
                                        )
                                        band_dos_mat.append(dos_matrix)
                                    kpt_dos_mat.append(band_dos_mat)
                                spin_dos_mat.append(kpt_dos_mat)
                        grand_dos_matrix = np.array(spin_dos_mat)
                        d["grand_dos_matrix"] = grand_dos_matrix
                        d["orbital_dict"] = orbital_dict

    def parse_scf(self, node):
        """
        Parses the total energy and dipole moments for a VASP calculation

        Args:
            node: (xml.etree.Element instance): The node to parse

        Returns:
            d (dict): Dictionary to containing parsed data
        """
        d = dict()
        if not (node.tag == "scstep"):
            raise AssertionError()
        for item in node:
            if item.tag == "energy":
                for i in item:
                    if i.attrib["name"] == "e_wo_entrp":
                        d["scf_energy"] = get_float_with_exception(i.text)
                    if i.attrib["name"] == "e_fr_energy":
                        d["scf_fr_energy"] = get_float_with_exception(i.text)
                    if i.attrib["name"] == "e_0_energy":
                        d["scf_0_energy"] = get_float_with_exception(i.text)
            if item.tag == "dipole":
                for i in item:
                    if i.attrib["name"] == "dipole":
                        d["scf_dipole_moment"] = self._parse_vector(i, vec_type=float)
        return d

    def parse_calc_to_dict(self, node, d):
        """
        Parses ionic step data from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        scf_energies = list()
        scf_fr_energies = list()
        scf_0_energies = list()
        scf_moments = list()
        for item in node:
            if item.tag in ["scstep"]:
                scf_dict = self.parse_scf(item)
                scf_energies.append(scf_dict["scf_energy"])
                scf_fr_energies.append(scf_dict["scf_fr_energy"])
                scf_0_energies.append(scf_dict["scf_0_energy"])
                if "scf_dipole_moment" in scf_dict.keys():
                    scf_moments.append(scf_dict["scf_dipole_moment"])
            if item.tag in ["structure"]:
                struct_dict = dict()
                self.parse_structure_to_dict(item, struct_dict)
                d["positions"].append(struct_dict["positions"])
                d["cells"].append(struct_dict["cell"])
            if item.tag in ["varray"] and item.attrib["name"] == "forces":
                d["forces"].append(self._parse_2d_matrix(item, vec_type=float))
            if item.tag in ["stress"]:
                d["stress_tensors"].append(self._parse_2d_matrix(item, vec_type=float))
            if item.tag == "energy":
                for i in item:
                    if i.attrib["name"] == "e_wo_entrp":
                        d["total_energies"].append(float(i.text))
                    if i.attrib["name"] == "e_fr_energy":
                        d["total_fr_energies"].append(float(i.text))
                    if i.attrib["name"] == "e_0_energy":
                        d["total_0_energies"].append(float(i.text))
                    if i.attrib["name"] == "kinetic":
                        d["kinetic_energies"].append(float(i.text))
            if item.tag == "eigenvalues":
                self.parse_eigenvalues_to_dict(item, d)

            if item.tag == "dos":
                self.parse_fermi_level_to_dict(item, d)
                d["efermi"] = float(d["efermi"])
                for i in item:
                    if i.tag == "total":
                        try:
                            self.parse_total_dos_to_dict(i, d)
                        except ValueError:
                            pass
                    if i.tag == "partial":
                        try:
                            self.parse_partial_dos_to_dict(i, d)
                        except ValueError:
                            pass

            if item.tag == "projected":
                self.parse_projected_dos_to_dict(item, d)

            if "cce" in item.tag:
                self.parse_cce_to_dict(item, d)

        d["scf_energies"].append(scf_energies)
        d["scf_fr_energies"].append(scf_fr_energies)
        d["scf_0_energies"].append(scf_0_energies)
        d["scf_dipole_moments"].append(scf_moments)

    @staticmethod
    def parse_cce_to_dict(node, d):
        for item in node:
            if item.attrib["name"] not in list(d.keys()):
                d[item.attrib["name"]] = list()
            d[item.attrib["name"]].append(float(item.text))

    def parse_eigenvalues_to_dict(self, node, d):
        """
        Parses eigenvalue and occupancy data from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "eigenvalues"):
            raise AssertionError()
        grand_eigenvalue_matrix = list()
        grand_occupancy_matrix = list()
        for item in node:
            if item.tag == "array":
                for ii in item:
                    if ii.tag == "set":
                        spin_occ_mat = list()
                        spin_eig_mat = list()
                        for sp in ii:
                            if sp.tag == "set" and "spin" in sp.attrib["comment"]:
                                kpt_eig_mat = list()
                                kpt_occ_mat = list()
                                for kpt in sp:
                                    values = self._parse_2d_matrix(kpt, vec_type=float)
                                    eig_vec = values[:, 0].flatten()
                                    occ_vec = values[:, 1].flatten()
                                    kpt_eig_mat.append(eig_vec)
                                    kpt_occ_mat.append(occ_vec)
                                spin_eig_mat.append(kpt_eig_mat)
                                spin_occ_mat.append(kpt_occ_mat)
                        grand_eigenvalue_matrix = np.array(spin_eig_mat)
                        grand_occupancy_matrix = np.array(spin_occ_mat)
        d["grand_eigenvalue_matrix"] = grand_eigenvalue_matrix
        d["grand_occupancy_matrix"] = grand_occupancy_matrix

    def parse_structure_to_dict(self, node, d):
        """
        Parses structure from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "structure"):
            raise AssertionError()
        for leaf in node:
            if leaf.tag == "crystal":
                for item in leaf:
                    if item.tag == "varray" and item.attrib["name"] == "basis":
                        d["cell"] = self._parse_2d_matrix(item)
            if leaf.tag == "varray" and leaf.attrib["name"] == "positions":
                d["positions"] = self._parse_2d_matrix(leaf)

            if leaf.tag == "varray" and leaf.attrib["name"] == "selective":
                d["selective_dynamics"] = self._parse_2d_matrix(leaf, vec_type=bool)

    @staticmethod
    def parse_item_to_dict(node, d):
        """
        Parses values from an item to a dictionary

        Args:
            node (etree.Element instance): Node to be parsed
            d (dict): The dictionary to which data is to be parsed

        Returns:
            d (dictionary)
        """
        type_dict = {"string": str, "float": float, "int": int, "logical": bool}
        logical_dict = {"T": True, "F": False}
        try:
            if node.attrib["type"] == "logical":
                d[node.attrib["name"]] = logical_dict[node.text.strip()]
            else:
                d[node.attrib["name"]] = type_dict[node.attrib["type"]](node.text)
        except (KeyError, IndexError, ValueError):
            d[node.attrib["name"]] = node.text
        return d

    def parse_parameters(self, node, d):
        """
        Parses parameter data from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
        """
        if not (node.tag == "parameters"):
            raise AssertionError()
        self.parse_recursively(node, d, key_name="parameters")

    def parse_recursively(self, node, d, key_name=None):
        """
        Parses recursively from a node to a dictionary

        Args:
            node (xml.etree.Element instance): The node to parse
            d (dict): The dictionary to which data is to be parsed
            key_name (str): Forcefully assign a key name in case it is not present in the xml file
        """
        if not len(node) > 0:
            d[clean_key(node.attrib["name"])] = clean_character(node.text)
            return
        else:
            try:
                if key_name is not None:
                    dict_key = clean_key(key_name)
                else:
                    dict_key = clean_key(node.attrib["name"])
                d[dict_key] = dict()
                for item in node:
                    try:
                        self.parse_item_to_dict(item, d[dict_key][item.attrib["name"]])
                    except (KeyError, ValueError, IndexError):
                        try:
                            self.parse_recursively(
                                item, d[dict_key], item.attrib["name"]
                            )
                        except (KeyError, ValueError, IndexError):
                            pass
            except KeyError:
                pass

    def _parse_2d_matrix(self, node, vec_type=float):
        """
        Parses a 2D vector from a node

        Args:
            node (xml.etree.Element instance): The node to parse
            vec_type (type): The type of the vector to be parsed

        Returns:
            numpy.ndarray: The required 2D array/vector
        """
        arr = list()
        for item in node:
            arr.append(self._parse_vector(item, vec_type=vec_type))
        return np.array(arr)

    @staticmethod
    def _parse_vector(node, vec_type=float):
        """
        Parses a 1D vector from a node

        Args:
            node (xml.etree.Element instance): The node to parse
            vec_type (type): The type of the vector to be parsed

        Returns:
            numpy.ndarray: The required 1D array/vector
        """
        txt = node.text
        lst = txt.split()
        logical_dict = {"T": True, "F": False}
        if "type" in node.attrib.keys():
            if node.attrib["type"] == "logical":
                return np.array([logical_dict[l.strip()] for l in lst])
            else:
                return np.array([vec_type(l) for l in lst])
        else:
            return np.array([vec_type(l) for l in lst])

    def get_initial_structure(self):
        """
        Gets the initial structure from the simulation

        Returns:
            pyiron.atomistics.structure.atoms.Atoms: The initial structure

        """
        # try:
        el_list = [
            re.split("[^a-zA-Z]", el)[0]
            for el in self.vasprun_dict["atominfo"]["species_list"]
        ]
        cell = self.vasprun_dict["init_structure"]["cell"]
        positions = self.vasprun_dict["init_structure"]["positions"]
        if len(positions[positions > 1.01]) > 0:
            basis = Atoms(el_list, positions=positions, cell=cell, pbc=True)
        else:
            basis = Atoms(el_list, scaled_positions=positions, cell=cell, pbc=True)
        if "selective_dynamics" in self.vasprun_dict["init_structure"].keys():
            constraints_dict = {
                label: []
                for label in ["TTT", "TTF", "FTT", "TFT", "TFF", "FFT", "FTF", "FFF"]
            }
            for i, val in enumerate(
                self.vasprun_dict["init_structure"]["selective_dynamics"]
            ):
                if val[0] and val[1] and val[2]:
                    constraints_dict["TTT"].append(i)
                elif val[0] and val[1] and not val[2]:
                    constraints_dict["TTF"].append(i)
                elif not val[0] and val[1] and val[2]:
                    constraints_dict["FTT"].append(i)
                elif val[0] and not val[1] and val[2]:
                    constraints_dict["TFT"].append(i)
                elif val[0] and not val[1] and not val[2]:
                    constraints_dict["TFF"].append(i)
                elif not val[0] and val[1] and not val[2]:
                    constraints_dict["FTF"].append(i)
                elif not val[0] and not val[1] and val[2]:
                    constraints_dict["FFT"].append(i)
                elif not val[0] and not val[1] and not val[2]:
                    constraints_dict["FFF"].append(i)
                else:
                    raise ValueError("Selective Dynamics Error: " + str(val))

            constraints_lst = []
            for k, v in constraints_dict.items():
                if len(v) > 0:
                    if k == "TTT":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(False, False, False))
                        )
                    elif k == "TTF":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(False, False, True))
                        )
                    elif k == "FTT":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(True, False, False))
                        )
                    elif k == "TFT":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(False, True, False))
                        )
                    elif k == "TFF":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(False, True, True))
                        )
                    elif k == "FTF":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(True, False, True))
                        )
                    elif k == "FFT":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(True, True, False))
                        )
                    elif k == "FFF":
                        constraints_lst.append(
                            FixCartesian(a=v, mask=(True, True, True))
                        )
                    else:
                        raise ValueError(
                            "Selective Dynamics Error: " + str(k) + ": " + str(v)
                        )

            basis.set_constraint(constraints_lst)
        return basis

    def get_final_structure(self):
        """
        Gets the final structure from the simulation

        Returns:
            pyiron.atomistics.structure.atoms.Atoms: The final structure

        """
        # try:
        basis = self.get_initial_structure()
        basis.set_cell(self.vasprun_dict["final_structure"]["cell"])
        positions = self.vasprun_dict["final_structure"]["positions"]
        if len(positions[positions > 1.01]) > 0:
            basis.positions = positions
        else:
            basis.set_scaled_positions(positions)
        return basis
        # except (KeyError, AttributeError, ValueError):
        #     return

    def get_electronic_structure(self, es_class=ElectronicStructure):
        """
        Get's the electronic structure from the VASP calculation

        Returns:
            pyiron.atomistics.waves.electronic.ElectronicStructure: The electronic structure object

        """
        es_obj = es_class()
        es_obj.kpoint_list = self.vasprun_dict["kpoints"]["kpoint_list"]
        es_obj.kpoint_weights = self.vasprun_dict["kpoints"]["kpoint_weights"]
        es_obj.eigenvalue_matrix = self.vasprun_dict["grand_eigenvalue_matrix"]
        es_obj.occupancy_matrix = self.vasprun_dict["grand_occupancy_matrix"]
        es_obj.n_spins = len(es_obj.occupancy_matrix)
        if "grand_dos_matrix" in self.vasprun_dict.keys():
            es_obj.grand_dos_matrix = self.vasprun_dict["grand_dos_matrix"]
        if "efermi" in self.vasprun_dict.keys():
            es_obj.efermi = self.vasprun_dict["efermi"]
        if "spin_dos_energies" in self.vasprun_dict.keys():
            es_obj.dos_energies = self.vasprun_dict["spin_dos_energies"][0]
            es_obj.dos_densities = self.vasprun_dict["spin_dos_density"]
            es_obj.dos_idensities = self.vasprun_dict["spin_dos_idensity"]
        if "resolved_dos_matrix" in self.vasprun_dict.keys():
            es_obj.resolved_densities = self.vasprun_dict["resolved_dos_matrix"]
            es_obj.orbital_dict = self.vasprun_dict["orbital_dict"]
        es_obj.generate_from_matrices()
        return es_obj

    def get_potentiostat_output(self):
        """
        Gets the input from the vasp potentiostat output and stores it using intuitive names

        Returns:
            dict/None: Dictionary with potentiostat output

        """
        potstat_dict = dict()
        if "dftnw_pot" not in self.vasprun_dict.keys():
            return
        try:
            vasprun_dict_keys = [
                "dftnw_pot",
                "dftnw_zval",
                "dftnw_electrodecharge",
                "dftnw_vaclevel_upperside_surf",
                "dftnw_vaclevel_lowerside_surf",
                "dftnw_efermi",
            ]
            potstat_dict_keys = [
                "potential_drop",
                "Ne_charge",
                "electrode_charge",
                "vac_level_upper",
                "vac_level_lower",
                "fermi_level",
            ]
            for k, vk in zip(potstat_dict_keys, vasprun_dict_keys):
                if vk in self.vasprun_dict.keys():
                    potstat_dict[k] = np.array(self.vasprun_dict[vk])
            return potstat_dict
        except KeyError:
            if len(potstat_dict.keys()) > 0:
                return potstat_dict
            else:
                return

    def get_valence_electrons_per_atom(self):
        """
        Gets the number of valence electrons for each atom in the final structure

        Returns:
            numpy.ndarray: Array of valence charges
        """
        sp_dict = self.vasprun_dict["atominfo"]["species_dict"]
        return np.hstack(
            [[val["valence"]] * val["n_atoms"] for val in sp_dict.values()]
        )


def clean_character(a, remove_char=" "):
    """
    Args:
        a (str): String to be cleaned
        remove_char (str): Character to be replaced

    Returns:
        str: The clean string
    """
    if isinstance(a, str):
        return a.replace(remove_char, "")
    else:
        return a


def clean_key(a, remove_char=" "):
    """
    Replaces blanck spaces from a string for a dictionary key with "_"

    Args:
        a (str): String to be cleaned
        remove_char (str): Character to be replaced

    Returns:
        str: The clean string
    """
    if isinstance(a, str):
        return a.replace(remove_char, "_")
    else:
        return a


def get_float_with_exception(text, exception_value=0.0):
    """
    Converts a text into a corresponding float or returns `exception_value` if it can't do this

    Args:
        text (str/numpy.str_): String to convert to float
        exception_value (float/None): Value to be returned if you can't text to a float

    Returns:
        float/None: Exception value

    """
    try:
        return float(text)
    except ValueError:
        warnings.warn(message=" ", category=VasprunWarning)
        return exception_value


class VasprunError(ValueError):
    pass


class VasprunWarning(UserWarning):
    pass
