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


"""
This module provides a class used to describe the elastic tensor,
including methods used to fit the elastic tensor from linear response
stress-strain data
"""

import itertools
import warnings
from collections import OrderedDict

import numpy as np
import sympy as sp
from monty.dev import deprecated
from scipy.integrate import quad
from scipy.optimize import root
from scipy.special import factorial

from pymatgen.analysis.elasticity.strain import Strain
from pymatgen.analysis.elasticity.stress import Stress
from pymatgen.core.tensors import (
    DEFAULT_QUAD,
    SquareTensor,
    Tensor,
    TensorCollection,
    get_uvec,
)
from pymatgen.core.units import Unit

__author__ = "Joseph Montoya"
__copyright__ = "Copyright 2012, The Materials Project"
__credits__ = "Maarten de Jong, Ian Winter, Shyam Dwaraknath, " "Mark Asta, Anubhav Jain"
__version__ = "1.0"
__maintainer__ = "Joseph Montoya"
__email__ = "montoyjh@lbl.gov"
__status__ = "Production"
__date__ = "July 24, 2018"


class NthOrderElasticTensor(Tensor):
    """
    An object representing an nth-order tensor expansion
    of the stress-strain constitutive equations
    """

    GPa_to_eV_A3 = Unit("GPa").get_conversion_factor(Unit("eV ang^-3"))
    symbol = "C"

    def __new__(cls, input_array, check_rank=None, tol=1e-4):
        """
        Args:
            input_array ():
            check_rank ():
            tol ():
        """
        obj = super().__new__(cls, input_array, check_rank=check_rank)
        if obj.rank % 2 != 0:
            raise ValueError("ElasticTensor must have even rank")
        if not obj.is_voigt_symmetric(tol):
            warnings.warn("Input elastic tensor does not satisfy " "standard voigt symmetries")
        return obj.view(cls)

    @property
    def order(self):
        """
        Order of the elastic tensor
        """
        return self.rank // 2

    def calculate_stress(self, strain):
        """
        Calculate's a given elastic tensor's contribution to the
        stress using Einstein summation

        Args:
            strain (3x3 array-like): matrix corresponding to strain
        """
        strain = np.array(strain)
        if strain.shape == (6,):
            strain = Strain.from_voigt(strain)
        assert strain.shape == (3, 3), "Strain must be 3x3 or voigt-notation"
        stress_matrix = self.einsum_sequence([strain] * (self.order - 1)) / factorial(self.order - 1)
        return Stress(stress_matrix)

    def energy_density(self, strain, convert_GPa_to_eV=True):
        """
        Calculates the elastic energy density due to a strain
        """
        e_density = np.sum(self.calculate_stress(strain) * strain) / self.order
        if convert_GPa_to_eV:
            e_density *= self.GPa_to_eV_A3  # Conversion factor for GPa to eV/A^3
        return e_density

    @classmethod
    def from_diff_fit(cls, strains, stresses, eq_stress=None, order=2, tol=1e-10):
        """

        Args:
            strains ():
            stresses ():
            eq_stress ():
            order ():
            tol ():

        Returns:

        """
        return cls(diff_fit(strains, stresses, eq_stress, order, tol)[order - 2])


def raise_error_if_unphysical(f):
    """
    Wrapper for functions or properties that should raise an error
    if tensor is unphysical.
    """

    def wrapper(self, *args, **kwargs):
        """
        Args:
            self ():
            *args ():
            **kwargs ():

        Returns:

        """
        if self.k_vrh < 0 or self.g_vrh < 0:
            raise ValueError("Bulk or shear modulus is negative, property " "cannot be determined")
        return f(self, *args, **kwargs)

    return wrapper


class ElasticTensor(NthOrderElasticTensor):
    """
    This class extends Tensor to describe the 3x3x3x3
    second-order elastic tensor, C_{ijkl}, with various
    methods for estimating other properties derived from
    the second order elastic tensor
    """

    def __new__(cls, input_array, tol=1e-4):
        """
        Create an ElasticTensor object.  The constructor throws an error if
        the shape of the input_matrix argument is not 3x3x3x3, i. e. in true
        tensor notation.  Issues a warning if the input_matrix argument does
        not satisfy standard symmetries.  Note that the constructor uses
        __new__ rather than __init__ according to the standard method of
        subclassing numpy ndarrays.

        Args:
            input_array (3x3x3x3 array-like): the 3x3x3x3 array-like
                representing the elastic tensor

            tol (float): tolerance for initial symmetry test of tensor
        """

        obj = super().__new__(cls, input_array, check_rank=4, tol=tol)
        return obj.view(cls)

    @property
    def compliance_tensor(self):
        """
        returns the Voigt-notation compliance tensor,
        which is the matrix inverse of the
        Voigt-notation elastic tensor
        """
        s_voigt = np.linalg.inv(self.voigt)
        return ComplianceTensor.from_voigt(s_voigt)

    @property
    def k_voigt(self):
        """
        returns the K_v bulk modulus
        """
        return self.voigt[:3, :3].mean()

    @property
    def g_voigt(self):
        """
        returns the G_v shear modulus
        """
        return (
            2.0 * self.voigt[:3, :3].trace() - np.triu(self.voigt[:3, :3]).sum() + 3 * self.voigt[3:, 3:].trace()
        ) / 15.0

    @property
    def k_reuss(self):
        """
        returns the K_r bulk modulus
        """
        return 1.0 / self.compliance_tensor.voigt[:3, :3].sum()

    @property
    def g_reuss(self):
        """
        returns the G_r shear modulus
        """
        return 15.0 / (
            8.0 * self.compliance_tensor.voigt[:3, :3].trace()
            - 4.0 * np.triu(self.compliance_tensor.voigt[:3, :3]).sum()
            + 3.0 * self.compliance_tensor.voigt[3:, 3:].trace()
        )

    @property
    def k_vrh(self):
        """
        returns the K_vrh (Voigt-Reuss-Hill) average bulk modulus
        """
        return 0.5 * (self.k_voigt + self.k_reuss)

    @property
    def g_vrh(self):
        """
        returns the G_vrh (Voigt-Reuss-Hill) average shear modulus
        """
        return 0.5 * (self.g_voigt + self.g_reuss)

    @property
    def y_mod(self):
        """
        Calculates Young's modulus (in SI units) using the
        Voigt-Reuss-Hill averages of bulk and shear moduli
        """
        return 9.0e9 * self.k_vrh * self.g_vrh / (3.0 * self.k_vrh + self.g_vrh)

    def directional_poisson_ratio(self, n, m, tol=1e-8):
        """
        Calculates the poisson ratio for a specific direction
        relative to a second, orthogonal direction

        Args:
            n (3-d vector): principal direction
            m (3-d vector): secondary direction orthogonal to n
            tol (float): tolerance for testing of orthogonality
        """
        n, m = get_uvec(n), get_uvec(m)
        if not np.abs(np.dot(n, m)) < tol:
            raise ValueError("n and m must be orthogonal")
        v = self.compliance_tensor.einsum_sequence([n] * 2 + [m] * 2)
        v *= -1 / self.compliance_tensor.einsum_sequence([n] * 4)
        return v

    def directional_elastic_mod(self, n):
        """
        Calculates directional elastic modulus for a specific vector
        """
        n = get_uvec(n)
        return self.einsum_sequence([n] * 4)

    @raise_error_if_unphysical
    def trans_v(self, structure):
        """
        Calculates transverse sound velocity (in SI units) using the
        Voigt-Reuss-Hill average bulk modulus

        Args:
            structure: pymatgen structure object

        Returns: transverse sound velocity (in SI units)

        """
        nsites = structure.num_sites
        volume = structure.volume
        natoms = structure.composition.num_atoms
        weight = float(structure.composition.weight)
        mass_density = 1.6605e3 * nsites * weight / (natoms * volume)
        if self.g_vrh < 0:
            raise ValueError("k_vrh or g_vrh is negative, " "sound velocity is undefined")
        return (1e9 * self.g_vrh / mass_density) ** 0.5

    @raise_error_if_unphysical
    def long_v(self, structure):
        """
        Calculates longitudinal sound velocity (in SI units)
        using the Voigt-Reuss-Hill average bulk modulus

        Args:
            structure: pymatgen structure object

        Returns: longitudinal sound velocity (in SI units)

        """
        nsites = structure.num_sites
        volume = structure.volume
        natoms = structure.composition.num_atoms
        weight = float(structure.composition.weight)
        mass_density = 1.6605e3 * nsites * weight / (natoms * volume)
        if self.g_vrh < 0:
            raise ValueError("k_vrh or g_vrh is negative, " "sound velocity is undefined")
        return (1e9 * (self.k_vrh + 4.0 / 3.0 * self.g_vrh) / mass_density) ** 0.5

    @raise_error_if_unphysical
    def snyder_ac(self, structure):
        """
        Calculates Snyder's acoustic sound velocity (in SI units)

        Args:
            structure: pymatgen structure object

        Returns: Snyder's acoustic sound velocity (in SI units)

        """
        nsites = structure.num_sites
        volume = structure.volume
        natoms = structure.composition.num_atoms
        num_density = 1e30 * nsites / volume
        tot_mass = sum([e.atomic_mass for e in structure.species])
        avg_mass = 1.6605e-27 * tot_mass / natoms
        return (
            0.38483
            * avg_mass
            * ((self.long_v(structure) + 2.0 * self.trans_v(structure)) / 3.0) ** 3.0
            / (300.0 * num_density ** (-2.0 / 3.0) * nsites ** (1.0 / 3.0))
        )

    @raise_error_if_unphysical
    def snyder_opt(self, structure):
        """
        Calculates Snyder's optical sound velocity (in SI units)

        Args:
            structure: pymatgen structure object

        Returns: Snyder's optical sound velocity (in SI units)

        """
        nsites = structure.num_sites
        volume = structure.volume
        num_density = 1e30 * nsites / volume
        return (
            1.66914e-23
            * (self.long_v(structure) + 2.0 * self.trans_v(structure))
            / 3.0
            / num_density ** (-2.0 / 3.0)
            * (1 - nsites ** (-1.0 / 3.0))
        )

    @raise_error_if_unphysical
    def snyder_total(self, structure):
        """
        Calculates Snyder's total sound velocity (in SI units)

        Args:
            structure: pymatgen structure object

        Returns: Snyder's total sound velocity (in SI units)

        """
        return self.snyder_ac(structure) + self.snyder_opt(structure)

    @raise_error_if_unphysical
    def clarke_thermalcond(self, structure):
        """
        Calculates Clarke's thermal conductivity (in SI units)

        Args:
            structure: pymatgen structure object

        Returns: Clarke's thermal conductivity (in SI units)

        """
        nsites = structure.num_sites
        volume = structure.volume
        tot_mass = sum([e.atomic_mass for e in structure.species])
        natoms = structure.composition.num_atoms
        weight = float(structure.composition.weight)
        avg_mass = 1.6605e-27 * tot_mass / natoms
        mass_density = 1.6605e3 * nsites * weight / (natoms * volume)
        return 0.87 * 1.3806e-23 * avg_mass ** (-2.0 / 3.0) * mass_density ** (1.0 / 6.0) * self.y_mod ** 0.5

    @raise_error_if_unphysical
    def cahill_thermalcond(self, structure):
        """
        Calculates Cahill's thermal conductivity (in SI units)

        Args:
            structure: pymatgen structure object

        Returns: Cahill's thermal conductivity (in SI units)

        """
        nsites = structure.num_sites
        volume = structure.volume
        num_density = 1e30 * nsites / volume
        return 1.3806e-23 / 2.48 * num_density ** (2.0 / 3.0) * (self.long_v(structure) + 2 * self.trans_v(structure))

    @raise_error_if_unphysical
    def debye_temperature(self, structure):
        """
        Estimates the debye temperature from longitudinal and
        transverse sound velocities

        Args:
            structure: pymatgen structure object

        Returns: debye temperature (in SI units)

        """
        v0 = structure.volume * 1e-30 / structure.num_sites
        vl, vt = self.long_v(structure), self.trans_v(structure)
        vm = 3 ** (1.0 / 3.0) * (1 / vl ** 3 + 2 / vt ** 3) ** (-1.0 / 3.0)
        td = 1.05457e-34 / 1.38065e-23 * vm * (6 * np.pi ** 2 / v0) ** (1.0 / 3.0)
        return td

    @deprecated(
        "debye_temperature_from_sound_velocities is now the default"
        "debye_temperature function, this one will be removed."
    )
    @raise_error_if_unphysical
    def debye_temperature_from_sound_velocities(self, structure):
        """
        Estimates Debye temperature from sound velocities
        """
        return self.debye_temperature(structure)

    @property
    def universal_anisotropy(self):
        """
        returns the universal anisotropy value
        """
        return 5.0 * self.g_voigt / self.g_reuss + self.k_voigt / self.k_reuss - 6.0

    @property
    def homogeneous_poisson(self):
        """
        returns the homogeneous poisson ratio
        """
        return (1.0 - 2.0 / 3.0 * self.g_vrh / self.k_vrh) / (2.0 + 2.0 / 3.0 * self.g_vrh / self.k_vrh)

    def green_kristoffel(self, u):
        """
        Returns the Green-Kristoffel tensor for a second-order tensor
        """
        return self.einsum_sequence([u, u], "ijkl,i,l")

    @property
    def property_dict(self):
        """
        returns a dictionary of properties derived from the elastic tensor
        """
        props = [
            "k_voigt",
            "k_reuss",
            "k_vrh",
            "g_voigt",
            "g_reuss",
            "g_vrh",
            "universal_anisotropy",
            "homogeneous_poisson",
            "y_mod",
        ]
        return {prop: getattr(self, prop) for prop in props}

    def get_structure_property_dict(self, structure, include_base_props=True, ignore_errors=False):
        """
        returns a dictionary of properties derived from the elastic tensor
        and an associated structure

        Args:
            structure (Structure): structure object for which to calculate
                associated properties
            include_base_props (bool): whether to include base properties,
                like k_vrh, etc.
            ignore_errors (bool): if set to true, will set problem properties
                that depend on a physical tensor to None, defaults to False
        """
        s_props = [
            "trans_v",
            "long_v",
            "snyder_ac",
            "snyder_opt",
            "snyder_total",
            "clarke_thermalcond",
            "cahill_thermalcond",
            "debye_temperature",
        ]
        if ignore_errors and (self.k_vrh < 0 or self.g_vrh < 0):
            sp_dict = {prop: None for prop in s_props}
        else:
            sp_dict = {prop: getattr(self, prop)(structure) for prop in s_props}
        sp_dict["structure"] = structure
        if include_base_props:
            sp_dict.update(self.property_dict)
        return sp_dict

    @classmethod
    def from_pseudoinverse(cls, strains, stresses):
        """
        Class method to fit an elastic tensor from stress/strain
        data.  Method uses Moore-Penrose pseudoinverse to invert
        the s = C*e equation with elastic tensor, stress, and
        strain in voigt notation

        Args:
            stresses (Nx3x3 array-like): list or array of stresses
            strains (Nx3x3 array-like): list or array of strains
        """
        # convert the stress/strain to Nx6 arrays of voigt-notation
        warnings.warn(
            "Pseudoinverse fitting of Strain/Stress lists may yield "
            "questionable results from vasp data, use with caution."
        )
        stresses = np.array([Stress(stress).voigt for stress in stresses])
        with warnings.catch_warnings(record=True):
            strains = np.array([Strain(strain).voigt for strain in strains])

        voigt_fit = np.transpose(np.dot(np.linalg.pinv(strains), stresses))
        return cls.from_voigt(voigt_fit)

    @classmethod
    def from_independent_strains(cls, strains, stresses, eq_stress=None, vasp=False, tol=1e-10):
        """
        Constructs the elastic tensor least-squares fit of independent strains
        Args:
            strains (list of Strains): list of strain objects to fit
            stresses (list of Stresses): list of stress objects to use in fit
                corresponding to the list of strains
            eq_stress (Stress): equilibrium stress to use in fitting
            vasp (boolean): flag for whether the stress tensor should be
                converted based on vasp units/convention for stress
            tol (float): tolerance for removing near-zero elements of the
                resulting tensor
        """
        strain_states = [tuple(ss) for ss in np.eye(6)]
        ss_dict = get_strain_state_dict(strains, stresses, eq_stress=eq_stress)
        if not set(strain_states) <= set(ss_dict.keys()):
            raise ValueError("Missing independent strain states: " "{}".format(set(strain_states) - set(ss_dict)))
        if len(set(ss_dict.keys()) - set(strain_states)) > 0:
            warnings.warn("Extra strain states in strain-stress pairs " "are neglected in independent strain fitting")
        c_ij = np.zeros((6, 6))
        for i in range(6):
            istrains = ss_dict[strain_states[i]]["strains"]
            istresses = ss_dict[strain_states[i]]["stresses"]
            for j in range(6):
                c_ij[i, j] = np.polyfit(istrains[:, i], istresses[:, j], 1)[0]
        if vasp:
            c_ij *= -0.1  # Convert units/sign convention of vasp stress tensor
        c = cls.from_voigt(c_ij)
        c = c.zeroed(tol)
        return c


class ComplianceTensor(Tensor):
    """
    This class represents the compliance tensor, and exists
    primarily to keep the voigt-conversion scheme consistent
    since the compliance tensor has a unique vscale
    """

    def __new__(cls, s_array):
        """
        Args:
            s_array ():
        """
        vscale = np.ones((6, 6))
        vscale[3:] *= 2
        vscale[:, 3:] *= 2
        obj = super().__new__(cls, s_array, vscale=vscale)
        return obj.view(cls)


class ElasticTensorExpansion(TensorCollection):
    """
    This class is a sequence of elastic tensors corresponding
    to an elastic tensor expansion, which can be used to
    calculate stress and energy density and inherits all
    of the list-based properties of TensorCollection
    (e. g. symmetrization, voigt conversion, etc.)
    """

    def __init__(self, c_list):
        """
        Initialization method for ElasticTensorExpansion

        Args:
            c_list (list or tuple): sequence of Tensor inputs
                or tensors from which the elastic tensor
                expansion is constructed.
        """
        c_list = [NthOrderElasticTensor(c, check_rank=4 + i * 2) for i, c in enumerate(c_list)]
        super().__init__(c_list)

    @classmethod
    def from_diff_fit(cls, strains, stresses, eq_stress=None, tol=1e-10, order=3):
        """
        Generates an elastic tensor expansion via the fitting function
        defined below in diff_fit
        """
        c_list = diff_fit(strains, stresses, eq_stress, order, tol)
        return cls(c_list)

    @property
    def order(self):
        """
        Order of the elastic tensor expansion, i. e. the order of the
        highest included set of elastic constants
        """
        return self[-1].order

    def calculate_stress(self, strain):
        """
        Calculate's a given elastic tensor's contribution to the
        stress using Einstein summation
        """
        return sum([c.calculate_stress(strain) for c in self])

    def energy_density(self, strain, convert_GPa_to_eV=True):
        """
        Calculates the elastic energy density due to a strain
        """
        return sum([c.energy_density(strain, convert_GPa_to_eV) for c in self])

    def get_ggt(self, n, u):
        """
        Gets the Generalized Gruneisen tensor for a given
        third-order elastic tensor expansion.

        Args:
            n (3x1 array-like): normal mode direction
            u (3x1 array-like): polarization direction
        """
        gk = self[0].einsum_sequence([n, u, n, u])
        result = -(
            2 * gk * np.outer(u, u) + self[0].einsum_sequence([n, n]) + self[1].einsum_sequence([n, u, n, u])
        ) / (2 * gk)
        return result

    def get_tgt(self, temperature=None, structure=None, quad=None):
        """
        Gets the thermodynamic Gruneisen tensor (TGT) by via an
        integration of the GGT weighted by the directional heat
        capacity.

        See refs:
            R. N. Thurston and K. Brugger, Phys. Rev. 113, A1604 (1964).
            K. Brugger Phys. Rev. 137, A1826 (1965).

        Args:
            temperature (float): Temperature in kelvin, if not specified
                will return non-cv-normalized value
            structure (float): Structure to be used in directional heat
                capacity determination, only necessary if temperature
                is specified
            quad (dict): quadrature for integration, should be
                dictionary with "points" and "weights" keys defaults
                to quadpy.sphere.Lebedev(19) as read from file
        """
        if temperature and not structure:
            raise ValueError("If using temperature input, you must also " "include structure")

        quad = quad if quad else DEFAULT_QUAD
        points = quad["points"]
        weights = quad["weights"]
        num, denom, c = np.zeros((3, 3)), 0, 1
        for p, w in zip(points, weights):
            gk = ElasticTensor(self[0]).green_kristoffel(p)
            rho_wsquareds, us = np.linalg.eigh(gk)
            us = [u / np.linalg.norm(u) for u in np.transpose(us)]
            for u in us:
                # TODO: this should be benchmarked
                if temperature:
                    c = self.get_heat_capacity(temperature, structure, p, u)
                num += c * self.get_ggt(p, u) * w
                denom += c * w
        return SquareTensor(num / denom)

    def get_gruneisen_parameter(self, temperature=None, structure=None, quad=None):
        """
        Gets the single average gruneisen parameter from the TGT.

        Args:
            temperature (float): Temperature in kelvin, if not specified
                will return non-cv-normalized value
            structure (float): Structure to be used in directional heat
                capacity determination, only necessary if temperature
                is specified
            quad (dict): quadrature for integration, should be
                dictionary with "points" and "weights" keys defaults
                to quadpy.sphere.Lebedev(19) as read from file
        """
        return np.trace(self.get_tgt(temperature, structure, quad)) / 3.0

    def get_heat_capacity(self, temperature, structure, n, u, cutoff=1e2):
        """
        Gets the directional heat capacity for a higher order tensor
        expansion as a function of direction and polarization.

        Args:
            temperature (float): Temperature in kelvin
            structure (float): Structure to be used in directional heat
                capacity determination
            n (3x1 array-like): direction for Cv determination
            u (3x1 array-like): polarization direction, note that
                no attempt for verification of eigenvectors is made
            cutoff (float): cutoff for scale of kt / (hbar * omega)
                if lower than this value, returns 0
        """
        k = 1.38065e-23
        kt = k * temperature
        hbar_w = 1.05457e-34 * self.omega(structure, n, u)
        if hbar_w > kt * cutoff: