# Source code for pymatgen.entries.computed_entries

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

    """
    This module implements equivalents of the basic ComputedEntry objects, which
    is the basic entity that can be used to perform many analyses. ComputedEntries
    contain calculated information, typically from VASP or other electronic
    structure codes. For example, ComputedEntries can be used as inputs for phase
    diagram analysis.
    """

    import json
    import abc

    from monty.json import MontyEncoder, MontyDecoder, MSONable

    from pymatgen.core.composition import Composition
    from pymatgen.core.structure import Structure
    from pymatgen.entries import Entry

    __author__ = "Ryan Kingsbury, Shyue Ping Ong, Anubhav Jain"
    __copyright__ = "Copyright 2011-2020, The Materials Project"
    __version__ = "1.1"
    __maintainer__ = "Shyue Ping Ong"
    __email__ = "shyuep@gmail.com"
    __status__ = "Production"
    __date__ = "April 2020"

    class EnergyAdjustment(MSONable):
        """
        Lightweight class to contain information about an energy adjustment or
        energy correction.
        """
        def __init__(self, value, name="Manual adjustment", cls=None, description=""):
            """
            Args:
                value: float, value of the energy adjustment in eV
                name: str, human-readable name of the energy adjustment.
                    (Default: Manual adjustment)
                cls: dict, Serialized Compatibility class used to generate the energy adjustment. (Default: None)
                description: str, human-readable explanation of the energy adjustment.
            """
            self.name = name
            self.cls = cls if cls else {}
            self.description = description

        @property
        @abc.abstractmethod
        def value(self):
            """
            Return the value of the energy adjustment in eV
            """

        def __repr__(self):
            output = ["{}:".format(self.__class__.__name__),
                      "  Name: {}".format(self.name),
                      "  Value: {:.3f} eV".format(self.value),
                      "  Description: {}".format(self.description),
                      "  Generated by: {}".format(self.cls.get("@class", None))]
            return "\n".join(output)

        @abc.abstractmethod
        def _normalize(self, factor):
            """
            Scale the value of the energy adjustment by factor.

            This method is utilized in ComputedEntry.normalize() to scale the energies to a formula unit basis
            (e.g. E_Fe6O9 = 3 x E_Fe2O3).
            """

    class ConstantEnergyAdjustment(EnergyAdjustment):
        """
        A constant energy adjustment applied to a ComputedEntry. Useful in energy referencing
        schemes such as the Aqueous energy referencing scheme.
        """
        def __init__(self, value, name="Constant energy adjustment", cls=None, description="Constant energy adjustment"):
            """
            Args:
                value: float, value of the energy adjustment in eV
                name: str, human-readable name of the energy adjustment.
                    (Default: Constant energy adjustment)
                cls: dict, Serialized Compatibility class used to generate the energy adjustment. (Default: None)
                description: str, human-readable explanation of the energy adjustment.
            """
            description = description + " ({:.3f} eV)".format(value)
            super().__init__(value, name, cls, description)
            self._value = value

        @property
        def value(self):
            """
            Return the value of the energy correction in eV.
            """
            return self._value

        @value.setter
        def value(self, x):
            self._value = x

        def _normalize(self, factor):
            self._value /= factor

    class ManualEnergyAdjustment(ConstantEnergyAdjustment):
        """
        A manual energy adjustment applied to a ComputedEntry.
        """
        def __init__(self, value):
            """
            Args:
                value: float, value of the energy adjustment in eV
            """
            name = "Manual energy adjustment"
            description = "Manual energy adjustment"
            super().__init__(value, name, cls=None, description=description)

    class CompositionEnergyAdjustment(EnergyAdjustment):
        """
        An energy adjustment applied to a ComputedEntry based on the atomic composition.
        Used in various DFT energy correction schemes.
        """
        def __init__(self, adj_per_atom, n_atoms, name, cls=None, description="Composition-based energy adjustment"):
            """
            Args:
                adj_per_atom: float, energy adjustment to apply per atom, in eV/atom
                n_atoms: float or int, number of atoms
                name: str, human-readable name of the energy adjustment.
                    (Default: "")
                cls: dict, Serialized Compatibility class used to generate the energy adjustment. (Default: None)
                description: str, human-readable explanation of the energy adjustment.
            """
            self._value = adj_per_atom
            self.n_atoms = n_atoms
            self.cls = cls if cls else {}
            self.name = name
            self.description = description + " ({:.3f} eV/atom x {} atoms)".format(self._value,
                                                                                   self.n_atoms
                                                                                   )

        @property
        def value(self):
            """
            Return the value of the energy adjustment in eV.
            """
            return self._value * self.n_atoms

        def _normalize(self, factor):
            self.n_atoms /= factor

    class TemperatureEnergyAdjustment(EnergyAdjustment):
        """
        An energy adjustment applied to a ComputedEntry based on the temperature.
        Used, for example, to add entropy to DFT energies.
        """
        def __init__(self, adj_per_deg, temp, n_atoms, name="", cls=None,
                     description="Temperature-based energy adjustment"):
            """
            Args:
                adj_per_deg: float, energy adjustment to apply per degree K, in eV/atom
                temp: float, temperature in Kelvin
                n_atoms: float or int, number of atoms
                name: str, human-readable name of the energy adjustment.
                    (Default: "")
                cls: dict, Serialized Compatibility class used to generate the energy adjustment. (Default: None)
                description: str, human-readable explanation of the energy adjustment.
            """
            self._value = adj_per_deg
            self.temp = temp
            self.n_atoms = n_atoms
            self.name = name
            self.cls = cls if cls else {}
            self.description = description + " ({:.4f} eV/K/atom x {} K x {} atoms)".format(self._value,
                                                                                            self.temp,
                                                                                            self.n_atoms,
                                                                                            )

        @property
        def value(self):
            """
            Return the value of the energy correction in eV.
            """
            return self._value * self.temp * self.n_atoms

        def _normalize(self, factor):
            self.n_atoms /= factor

    class ComputedEntry(Entry):
        """
        Lightweight Entry object for computed data. Contains facilities
        for applying corrections to the .energy attribute and for storing
        calculation parameters.

        """

        def __init__(self,
                     composition: Composition,
                     energy: float,
                     correction: float = 0.0,
                     energy_adjustments: list = None,
                     parameters: dict = None,
                     data: dict = None,
                     entry_id: object = None):
            """
            Initializes a ComputedEntry.

            Args:
                composition (Composition): Composition of the entry. For
                    flexibility, this can take the form of all the typical input
                    taken by a Composition, including a {symbol: amt} dict,
                    a string formula, and others.
                energy (float): Energy of the entry. Usually the final calculated
                    energy from VASP or other electronic structure codes.
                energy_adjustments: An optional list of EnergyAdjustment to
                    be applied to the energy. This is used to modify the energy for
                    certain analyses. Defaults to None.
                parameters: An optional dict of parameters associated with
                    the entry. Defaults to None.
                data: An optional dict of any additional data associated
                    with the entry. Defaults to None.
                entry_id: An optional id to uniquely identify the entry.
            """
            super().__init__(composition, energy)
            self.uncorrected_energy = self._energy
            self.energy_adjustments = energy_adjustments if energy_adjustments else []

            if correction != 0.0:
                if energy_adjustments:
                    raise ValueError("Argument conflict! Setting correction = {:.3f} conflicts "
                                     "with setting energy_adjustments. Specify one or the "
                                     "other.".format(correction))

                self.correction = correction

            self.parameters = parameters if parameters else {}
            self.data = data if data else {}
            self.entry_id = entry_id
            self.name = self.composition.reduced_formula

        @property
        def energy(self) -> float:
            """
            :return: the *corrected* energy of the entry.
            """
            return self._energy + self.correction

        @property
        def correction(self) -> float:
            """
            Returns:
                float: the total energy correction / adjustment applied to the entry,
                    in eV.
            """
            return sum([e.value for e in self.energy_adjustments])

        @correction.setter
        def correction(self, x: float) -> None:
            corr = ManualEnergyAdjustment(x)
            self.energy_adjustments = [corr]

        def normalize(self, mode: str = "formula_unit") -> None:
            """
            Normalize the entry's composition and energy.

            Args:
                mode: "formula_unit" is the default, which normalizes to
                    composition.reduced_formula. The other option is "atom", which
                    normalizes such that the composition amounts sum to 1.
            """
            factor = self._normalization_factor(mode)
            self.uncorrected_energy /= factor
            for ea in self.energy_adjustments:
                ea._normalize(factor)
            super().normalize(mode)

        def __repr__(self):
            n_atoms = self.composition.num_atoms
            output = ["{} {:<10} - {:<12} ({})".format(str(self.entry_id),
                                                       type(self).__name__,
                                                       self.composition.formula,
                                                       self.composition.reduced_formula),
                      "{:<24} = {:<9.4f} eV ({:<8.4f} eV/atom)".format("Energy (Uncorrected)",
                                                                       self._energy,
                                                                       self._energy / n_atoms),
                      "{:<24} = {:<9.4f} eV ({:<8.4f} eV/atom)".format("Correction",
                                                                       self.correction,
                                                                       self.correction / n_atoms),
                      "{:<24} = {:<9.4f} eV ({:<8.4f} eV/atom)".format("Energy (Final)",
                                                                       self.energy,
                                                                       self.energy_per_atom),
                      "Energy Adjustments:"
                      ]
            if len(self.energy_adjustments) == 0:
                output.append("  None")
            else:
                for e in self.energy_adjustments:
                    output.append("  {:<23}: {:<9.4f} eV ({:<8.4f} eV/atom)".format(e.name,
                                                                                    e.value,
                                                                                    e.value / n_atoms))
            output.append("Parameters:")
            for k, v in self.parameters.items():
                output.append("  {:<22} = {}".format(k, v))
            output.append("Data:")
            for k, v in self.data.items():
                output.append("  {:<22} = {}".format(k, v))
            return "\n".join(output)

        def __str__(self):
            return self.__repr__()

        @classmethod
        def from_dict(cls, d) -> 'ComputedEntry':
            """
            :param d: Dict representation.
            :return: ComputedEntry
            """
            dec = MontyDecoder()
            # the first block here is for legacy ComputedEntry that were
            # serialized before we had the energy_adjustments attribute.
            if d["correction"] != 0 and not d.get("energy_adjustments"):
                return cls(d["composition"], d["energy"], d["correction"],
                           parameters={k: dec.process_decoded(v)
                                       for k, v in d.get("parameters", {}).items()},
                           data={k: dec.process_decoded(v)
                                 for k, v in d.get("data", {}).items()},
                           entry_id=d.get("entry_id", None))
            # this is the preferred / modern way of instantiating ComputedEntry
            # we don't pass correction explicitly because it will be calculated
            # on the fly from energy_adjustments
            else:
                return cls(d["composition"], d["energy"], correction=0,
                           energy_adjustments=[dec.process_decoded(e)
                                               for e in d.get("energy_adjustments", {})],
                           parameters={k: dec.process_decoded(v)
                                       for k, v in d.get("parameters", {}).items()},
                           data={k: dec.process_decoded(v)
                                 for k, v in d.get("data", {}).items()},
                           entry_id=d.get("entry_id", None))

        def as_dict(self) -> dict:
            """
            :return: MSONable dict.
            """
            return_dict = super().as_dict()
            return_dict.update({"energy_adjustments": json.loads(json.dumps(self.energy_adjustments, cls=MontyEncoder)),
                                "parameters": json.loads(json.dumps(self.parameters, cls=MontyEncoder)),
                                "data": json.loads(json.dumps(self.data, cls=MontyEncoder)),
                                "entry_id": self.entry_id,
                                "correction": self.correction})
            return return_dict

    class ComputedStructureEntry(ComputedEntry):
        """
        A heavier version of ComputedEntry which contains a structure as well. The
        structure is needed for some analyses.
        """

        def __init__(self,
                     structure: Structure,
                     energy: float,
                     correction: float = 0.0,
                     energy_adjustments: list = None,
                     parameters: dict = None,
                     data: dict = None,
                     entry_id: object = None):
            """
            Initializes a ComputedStructureEntry.

            Args:
                structure (Structure): The actual structure of an entry.
                energy (float): Energy of the entry. Usually the final calculated
                    energy from VASP or other electronic structure codes.
                energy_adjustments: An optional list of EnergyAdjustment to
                    be applied to the energy. This is used to modify the energy for
                    certain analyses. Defaults to None.
                parameters: An optional dict of parameters associated with
                    the entry. Defaults to None.
                data: An optional dict of any additional data associated
                    with the entry. Defaults to None.
                entry_id: An optional id to uniquely identify the entry.
            """
            super().__init__(
                structure.composition, energy, correction=correction, energy_adjustments=energy_adjustments,
                parameters=parameters, data=data, entry_id=entry_id)
            self.structure = structure

        def as_dict(self) -> dict:
            """
            :return: MSONAble dict.
            """
            d = super().as_dict()
            d["@module"] = self.__class__.__module__
            d["@class"] = self.__class__.__name__
            d["structure"] = self.structure.as_dict()
            return d

        @classmethod
        def from_dict(cls, d) -> 'ComputedStructureEntry':
            """
            :param d: Dict representation.
            :return: ComputedStructureEntry
            """
            dec = MontyDecoder()
            # the first block here is for legacy ComputedEntry that were
            # serialized before we had the energy_adjustments attribute.
            if d["correction"] != 0 and not d.get("energy_adjustments"):
                return cls(dec.process_decoded(d["structure"]), d["energy"], d["correction"],
                           parameters={k: dec.process_decoded(v)
                                       for k, v in d.get("parameters", {}).items()},
                           data={k: dec.process_decoded(v)
                                 for k, v in d.get("data", {}).items()},
                           entry_id=d.get("entry_id", None))
            # this is the preferred / modern way of instantiating ComputedEntry
            # we don't pass correction explicitly because it will be calculated
            # on the fly from energy_adjustments
            else:
                return cls(dec.process_decoded(d["structure"]), d["energy"], correction=0,
                           energy_adjustments=[dec.process_decoded(e)
                                               for e in d.get("energy_adjustments", {})],
                           parameters={k: dec.process_decoded(v)
                                       for k, v in d.get("parameters", {}).items()},
                           data={k: dec.process_decoded(v)
                                 for k, v in d.get("data", {}).items()},
                           entry_id=d.get("entry_id", None))