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

"""
This module provides class to generate and analyze interfacial reactions.
"""

import warnings
import numpy as np
import matplotlib.pylab as plt

from pymatgen import Composition
from pymatgen.analysis.phase_diagram import GrandPotentialPhaseDiagram
from pymatgen.analysis.reaction_calculator import Reaction

__author__ = "Yihan Xiao"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Yihan Xiao"
__email__ = "eric.xyh2011@gmail.com"
__status__ = "Production"
__date__ = "Aug 15 2017"


class InterfacialReactivity:
    """
    An object encompassing all relevant data for interface reactions.
    """
    EV_TO_KJ_PER_MOL = 96.4853

    def __init__(self, c1, c2, pd, norm=True, include_no_mixing_energy=False,
                 pd_non_grand=None, use_hull_energy=False):
        """
        Args:
            c1 (Composition): Composition object for reactant 1.
            c2 (Composition): Composition object for reactant 2.
            pd (PhaseDiagram): PhaseDiagram object or GrandPotentialPhaseDiagram
                object built from all elements in composition c1 and c2.
            norm (bool): Whether or not the total number of atoms in composition
                of reactant will be normalized to 1.
            include_no_mixing_energy (bool): No_mixing_energy for a reactant is the
                opposite number of its energy above grand potential convex hull. In
                cases where reactions involve elements reservoir, this param
                determines whether no_mixing_energy of reactants will be included
                in the final reaction energy calculation. By definition, if pd is
                not a GrandPotentialPhaseDiagram object, this param is False.
            pd_non_grand (PhaseDiagram): PhaseDiagram object but not
                GrandPotentialPhaseDiagram object built from elements in c1 and c2.
            use_hull_energy (bool): Whether or not use the convex hull energy for
                a given composition for reaction energy calculation. If false,
                the energy of ground state structure will be used instead.
                Note that in case when ground state can not be found for a
                composition, convex hull energy will be used associated with a
                warning message.
        """
        self.grand = isinstance(pd, GrandPotentialPhaseDiagram)

        # if include_no_mixing_energy is True, pd should be a
        # GrandPotentialPhaseDiagram object and pd_non_grand should be given.
        if include_no_mixing_energy and not self.grand:
            raise ValueError('Please provide grand phase diagram to compute'
                             ' no_mixing_energy!')
        if include_no_mixing_energy and not pd_non_grand:
            raise ValueError('Please provide non-grand phase diagram to '
                             'compute no_mixing_energy!')
        if self.grand and use_hull_energy and not pd_non_grand:
            raise ValueError('Please provide non-grand phase diagram if'
                             ' you want to use convex hull energy.')

        # Keeps copy of original compositions.
        self.c1_original = c1
        self.c2_original = c2

        # Two sets of composition attributes for two processing conditions:
        # normalization with and without exluding element(s) from reservoir.
        self.c1 = c1
        self.c2 = c2
        self.comp1 = c1
        self.comp2 = c2

        self.norm = norm
        self.pd = pd
        self.pd_non_grand = pd_non_grand
        self.use_hull_energy = use_hull_energy

        # Factor is the compositional ratio between composition self.c1 and
        # processed composition self.comp1. E.g., factor for
        # Composition('SiO2') and Composition('O') is 2.0. This factor will
        # be used to convert mixing ratio in self.comp1 - self.comp2
        # tie line to that in self.c1 - self.c2 tie line.
        self.factor1 = 1
        self.factor2 = 1

        if self.grand:
            # Excludes element(s) from reservoir.
            self.comp1 = Composition({k: v for k, v in c1.items()
                                      if k not in pd.chempots})
            self.comp2 = Composition({k: v for k, v in c2.items()
                                      if k not in pd.chempots})
            # Calculate the factors in case where self.grand = True and
            # self.norm = True.
            factor1 = self.comp1.num_atoms / c1.num_atoms
            factor2 = self.comp2.num_atoms / c2.num_atoms

        if self.norm:
            self.c1 = c1.fractional_composition
            self.c2 = c2.fractional_composition
            self.comp1 = self.comp1.fractional_composition
            self.comp2 = self.comp2.fractional_composition
            if self.grand:
                # Only when self.grand = True and self.norm = True
                # will self.factor be updated.
                self.factor1 = factor1
                self.factor2 = factor2

        # Computes energies for reactants in different scenarios.
        if not self.grand:
            if self.use_hull_energy:
                self.e1 = self.pd.get_hull_energy(self.comp1)
                self.e2 = self.pd.get_hull_energy(self.comp2)
            else:
                # Use entry energy as reactant energy if no reservoir
                # is present.
                self.e1 = InterfacialReactivity._get_entry_energy(
                    self.pd, self.comp1)
                self.e2 = InterfacialReactivity._get_entry_energy(
                    self.pd, self.comp2)
        else:
            if include_no_mixing_energy:
                # Computing grand potentials needs compositions containing
                # element(s) from reservoir, so self.c1 and self.c2 are used.
                self.e1 = self._get_grand_potential(self.c1)
                self.e2 = self._get_grand_potential(self.c2)
            else:
                self.e1 = self.pd.get_hull_energy(self.comp1)
                self.e2 = self.pd.get_hull_energy(self.comp2)

    @staticmethod
    def _get_entry_energy(pd, composition):
        """
        Finds the lowest entry energy for entries matching the composition.
        Entries with non-negative formation energies are excluded. If no
        entry is found, use the convex hull energy for the composition.

        Args:
            pd (PhaseDiagram): PhaseDiagram object.
            composition (Composition): Composition object that the target
            entry should match.

        Returns:
            The lowest entry energy among entries matching the composition.
        """
        candidate = [i.energy_per_atom for i in pd.qhull_entries if
                     i.composition.fractional_composition ==
                     composition.fractional_composition]

        if not candidate:
            warnings.warn("The reactant " + composition.reduced_formula +
                          " has no matching entry with negative formation"
                          " energy, instead convex hull energy for this"
                          " composition will be used for reaction energy "
                          "calculation. ")
            return pd.get_hull_energy(composition)
        else:
            min_entry_energy = min(candidate)
            return min_entry_energy * composition.num_atoms

    def _get_grand_potential(self, composition):
        """
        Computes the grand potential Phi at a given composition and
        chemical potential(s).

        Args:
            composition (Composition): Composition object.

        Returns:
            Grand potential at a given composition at chemical potential(s).
        """
        if self.use_hull_energy:
            grand_potential = self.pd_non_grand.get_hull_energy(composition)
        else:
            grand_potential = InterfacialReactivity._get_entry_energy(
                self.pd_non_grand, composition)
        grand_potential -= sum([composition[e] * mu
                                for e, mu in self.pd.chempots.items()])
        if self.norm:
            # Normalizes energy to the composition excluding element(s)
            # from reservoir.
            grand_potential /= sum([composition[el]
                                    for el in composition
                                    if el not in self.pd.chempots])
        return grand_potential

    def _get_energy(self, x):
        """
        Computes reaction energy in eV/atom at mixing ratio x : (1-x) for
        self.comp1 : self.comp2.

        Args:
            x (float): Mixing ratio x of reactants, a float between 0 and 1.

        Returns:
            Reaction energy.
        """
        return self.pd.get_hull_energy(self.comp1 * x + self.comp2 * (1 - x)) - self.e1 * x - self.e2 * (1 - x)

    def _get_reaction(self, x):
        """
        Generates balanced reaction at mixing ratio x : (1-x) for
        self.comp1 : self.comp2.

        Args:
            x (float): Mixing ratio x of reactants, a float between 0 and 1.

        Returns:
            Reaction object.
        """
        mix_comp = self.comp1 * x + self.comp2 * (1 - x)
        decomp = self.pd.get_decomposition(mix_comp)

        # Uses original composition for reactants.
        if np.isclose(x, 0):
            reactant = [self.c2_original]
        elif np.isclose(x, 1):
            reactant = [self.c1_original]
        else:
            reactant = list(set([self.c1_original, self.c2_original]))

        if self.grand:
            reactant += [Composition(e.symbol)
                         for e, v in self.pd.chempots.items()]

        product = [Composition(k.name) for k, v in decomp.items()]
        reaction = Reaction(reactant, product)

        x_original = self._get_original_composition_ratio(reaction)
        if np.isclose(x_original, 1):
            reaction.normalize_to(self.c1_original, x_original)
        else:
            reaction.normalize_to(self.c2_original, 1 - x_original)
        return reaction

    def _get_elmt_amt_in_rxt(self, rxt):
        """
        Computes total number of atoms in a reaction formula for elements
        not in external reservoir. This method is used in the calculation
        of reaction energy per mol of reaction formula.

        Args:
            rxt (Reaction): a reaction.

        Returns:
            Total number of atoms for non_reservoir elements.
        """
        return sum([rxt.get_el_amount(e) for e in self.pd.elements])

    def get_products(self):
        """
        List of formulas of potential products. E.g., ['Li','O2','Mn'].
        """
        products = set()
        for _, _, _, react, _ in self.get_kinks():
            products = products.union(set([k.reduced_formula
                                           for k in react.products]))
        return list(products)

    @staticmethod
    def _convert(x, factor1, factor2):
        """
        Converts mixing ratio x in comp1 - comp2 tie line to that in
        c1 - c2 tie line.

        Args:
            x (float): Mixing ratio x in comp1 - comp2 tie line, a float
                between 0 and 1.
            factor1 (float): Compositional ratio between composition c1 and
                processed composition comp1. E.g., factor for
                Composition('SiO2') and Composition('O') is 2.0.
            factor2 (float): Compositional ratio between composition c2 and
                processed composition comp2.

        Returns:
            Mixing ratio in c1 - c2 tie line, a float between 0 and 1.
        """
        return x * factor2 / ((1 - x) * factor1 + x * factor2)

    @staticmethod
    def _reverse_convert(x, factor1, factor2):
        """
        Converts mixing ratio x in c1 - c2 tie line to that in
        comp1 - comp2 tie line.

        Args:
            x (float): Mixing ratio x in c1 - c2 tie line, a float between
                0 and 1.
            factor1 (float): Compositional ratio between composition c1 and
                processed composition comp1. E.g., factor for
                Composition('SiO2') and Composition('O') is 2.
            factor2 (float): Compositional ratio between composition c2 and
                processed composition comp2.

        Returns:
            Mixing ratio in comp1 - comp2 tie line, a float between 0 and 1.
        """
        return x * factor1 / ((1 - x) * factor2 + x * factor1)

    def get_kinks(self):
        """
        Finds all the kinks in mixing ratio where reaction products changes
        along the tie line of composition self.c1 and composition self.c2.

        Returns:
            Zip object of tuples (index, mixing ratio,
                                  reaction energy per atom in eV/atom,
                                  reaction formula,
                                  reaction energy per mol of reaction
                                  formula in kJ/mol).
        """
        c1_coord = self.pd.pd_coords(self.comp1)
        c2_coord = self.pd.pd_coords(self.comp2)
        n1 = self.comp1.num_atoms
        n2 = self.comp2.num_atoms
        critical_comp = self.pd.get_critical_compositions(self.comp1,
                                                          self.comp2)
        x_kink, energy_kink, react_kink, energy_per_rxt_formula = \
            [], [], [], []
        if all(c1_coord == c2_coord):
            x_kink = [0, 1]
            energy_kink = [self._get_energy(x) for x in x_kink]
            react_kink = [self._get_reaction(x) for x in x_kink]
            num_atoms = [(x * self.comp1.num_atoms +
                          (1 - x) * self.comp2.num_atoms) for x in x_kink]
            energy_per_rxt_formula = [energy_kink[i] *
                                      self._get_elmt_amt_in_rxt(
                                          react_kink[i]) /
                                      num_atoms[i] *
                                      InterfacialReactivity.EV_TO_KJ_PER_MOL
                                      for i in range(2)]
        else:
            for i in reversed(critical_comp):
                # Gets mixing ratio x at kinks.
                c = self.pd.pd_coords(i)
                x = np.linalg.norm(c - c2_coord) / np.linalg.norm(c1_coord - c2_coord)
                # Modifies mixing ratio in case compositions self.comp1 and
                # self.comp2 are not normalized.
                x = x * n2 / (n1 + x * (n2 - n1))
                n_atoms = x * self.comp1.num_atoms + (1 - x) * self.comp2.num_atoms
                # Converts mixing ratio in comp1 - comp2 tie line to that in
                # c1 - c2 tie line.
                x_converted = InterfacialReactivity._convert(
                    x, self.factor1, self.factor2)
                x_kink.append(x_converted)
                # Gets reaction energy at kinks
                normalized_energy = self._get_energy(x)
                energy_kink.append(normalized_energy)
                # Gets balanced reaction at kinks
                rxt = self._get_reaction(x)
                react_kink.append(rxt)
                rxt_energy = normalized_energy * self._get_elmt_amt_in_rxt(rxt) / n_atoms
                energy_per_rxt_formula.append(
                    rxt_energy *
                    InterfacialReactivity.EV_TO_KJ_PER_MOL)
        index_kink = range(1, len(critical_comp) + 1)
        return zip(index_kink, x_kink, energy_kink, react_kink,
                   energy_per_rxt_formula)

    def get_critical_original_kink_ratio(self):
        """
        Returns a list of molar mixing ratio for each kink between ORIGINAL
        (instead of processed) reactant compositions. This is the
        same list as mixing ratio obtained from get_kinks method
        if self.norm = False.

        Returns:
            A list of floats representing molar mixing ratios between
            the original reactant compositions for each kink.
        """
        ratios = []
        if self.c1_original == self.c2_original:
            return [0, 1]
        reaction_kink = [k[3] for k in self.get_kinks()]
        for rxt in reaction_kink:
            ratios.append(abs(self._get_original_composition_ratio(rxt)))
        return ratios

    def _get_original_composition_ratio(self, reaction):
        """
        Returns the molar mixing ratio between the reactants with ORIGINAL (
        instead of processed) compositions for a reaction.

        Args:
            reaction (Reaction): Reaction object that contains the original
                reactant compositions.

        Returns:
            The molar mixing ratio between the original reactant
            compositions for a reaction.
        """
        if self.c1_original == self.c2_original:
            return 1
        c1_coeff = reaction.get_coeff(self.c1_original) \
            if self.c1_original in reaction.reactants else 0
        c2_coeff = reaction.get_coeff(self.c2_original) \
            if self.c2_original in reaction.reactants else 0
        return c1_coeff * 1.0 / (c1_coeff + c2_coeff)

    def labels(self):
        """
        Returns a dictionary containing kink information:
        {index: 'x= mixing_ratio energy= reaction_energy reaction_equation'}.
        E.g., {1: 'x= 0.0 energy = 0.0 Mn -> Mn',
               2: 'x= 0.5 energy = -15.0 O2 + Mn -> MnO2',
               3: 'x= 1.0 energy = 0.0 O2 -> O2'}.
        """
        return {j: 'x= ' + str(round(x, 4)) + ' energy in eV/atom = ' +
                   str(round(energy, 4)) + ' ' + str(reaction)
                for j, x, energy, reaction, _ in self.get_kinks()}

    def plot(self):
        """
        Plots reaction energy as a function of mixing ratio x in
        self.c1 - self.c2 tie line using pylab.

        Returns:
            Pylab object that plots reaction energy as a function of
            mixing ratio x.
        """
        plt.rcParams['xtick.major.pad'] = '6'
        plt.rcParams['ytick.major.pad'] = '6'
        plt.rcParams['axes.linewidth'] = 2
        npoint = 1000
        xs = np.linspace(0, 1, npoint)

        # Converts sampling points in self.c1 - self.c2 tie line to those in
        # self.comp1 - self.comp2 tie line.
        xs_reverse_converted = InterfacialReactivity._reverse_convert(
            xs, self.factor1, self.factor2)
        energies = [self._get_energy(x) for x in xs_reverse_converted]
        plt.plot(xs, energies, 'k-')

        # Marks kinks and minimum energy point.
        kinks = self.get_kinks()
        _, x_kink, energy_kink, _, _ = zip(*kinks)
        plt.scatter(x_kink, energy_kink, marker='o', c='blue', s=20)
        plt.scatter(self.minimum()[0], self.minimum()[1], marker='*',
                    c='red', s=300)

        # Labels kinks with indices. Labels are made draggable
        # in case of overlapping.
        for index, x, energy, _, _ in kinks:
            plt.annotate(
                index,
                xy=(x, energy), xytext=(5, 30),
                textcoords='offset points', ha='right', va='bottom',
                arrowprops=dict(arrowstyle='->',
                                connectionstyle='arc3,rad=0')).draggable()
        plt.xlim([-0.05, 1.05])
        if self.norm:
            plt.ylabel('Energy (eV/atom)')
        else:
            plt.ylabel('Energy (eV/f.u.)')
        plt.xlabel('$x$ in $x$ {} + $(1-x)$ {}'.format(
            self.c1.reduced_formula, self.c2.reduced_formula))
        return plt

    def minimum(self):
        """
        Finds the minimum reaction energy E_min and corresponding
        mixing ratio x_min.

        Returns:
            Tuple (x_min, E_min).
        """
        return min([(x, energy) for _, x, energy, _, _ in self.get_kinks()],
                   key=lambda i: i[1])

    def get_no_mixing_energy(self):
        """
        Generates the opposite number of energy above grand potential
        convex hull for both reactants.

        Returns:
            [(reactant1, no_mixing_energy1),(reactant2,no_mixing_energy2)].
        """
        assert self.grand == 1, 'Please provide grand potential phase diagram for computing no_mixing_energy!'

        energy1 = self.pd.get_hull_energy(self.comp1) - self._get_grand_potential(self.c1)
        energy2 = self.pd.get_hull_energy(self.comp2) - self._get_grand_potential(self.c2)
        unit = 'eV/f.u.'
        if self.norm:
            unit = 'eV/atom'
        return [(self.c1_original.reduced_formula +
                 ' ({0})'.format(unit), energy1),
                (self.c2_original.reduced_formula +
                 ' ({0})'.format(unit), energy2)]

    @staticmethod
    def get_chempot_correction(element, temp, pres):
        """
        Get the normalized correction term Δμ for chemical potential of a gas
        phase consisting of element at given temperature and pressure,
        referenced to that in the standard state (T_std = 298.15 K,
        T_std = 1 bar). The gas phase is limited to be one of O2, N2, Cl2,
        F2, H2. Calculation formula can be found in the documentation of
        Materials Project website.

        Args:
            element (string): The string representing the element.
            temp (float): The temperature of the gas phase.
            pres (float): The pressure of the gas phase.

        Returns:
            The correction of chemical potential in eV/atom of the gas
            phase at given temperature and pressure.
        """
        if element not in ["O", "N", "Cl", "F", "H"]:
            return 0
        std_temp = 298.15
        std_pres = 1E5
        ideal_gas_const = 8.3144598
        # Cp and S at standard state in J/(K.mol). Data from
        # https://janaf.nist.gov/tables/O-029.html
        # https://janaf.nist.gov/tables/N-023.html
        # https://janaf.nist.gov/tables/Cl-073.html
        # https://janaf.nist.gov/tables/F-054.html
        # https://janaf.nist.gov/tables/H-050.html
        Cp_dict = {"O": 29.376,
                   "N": 29.124,
                   "Cl": 33.949,
                   "F": 31.302,
                   "H": 28.836}

        S_dict = {"O": 205.147,
                  "N": 191.609,
                  "Cl": 223.079,
                  "F": 202.789,
                  "H": 130.680}
        Cp_std = Cp_dict[element]
        S_std = S_dict[element]
        PV_correction = ideal_gas_const * temp * np.log(pres / std_pres)
        TS_correction = - Cp_std * (temp * np.log(temp) - std_temp * np.log(std_temp)) \
            + Cp_std * (temp - std_temp) * (1 + np.log(std_temp)) \
            - S_std * (temp - std_temp)

        dG = PV_correction + TS_correction
        # Convert to eV/molecule unit.
        dG /= 1000 * InterfacialReactivity.EV_TO_KJ_PER_MOL
        # Normalize by number of atoms in the gas molecule. For elements
        # considered, the gas molecules are all diatomic.
        dG /= 2
        return dG
