from __future__ import annotations

import warnings
from copy import deepcopy

import numpy as np
import orjson
import pytest
from numpy.testing import assert_allclose
from pytest import approx
from scipy import linalg

from pymatgen.core.elasticity.elastic import (
    ComplianceTensor,
    ElasticTensor,
    ElasticTensorExpansion,
    NthOrderElasticTensor,
    diff_fit,
    find_eq_stress,
    generate_pseudo,
    get_diff_coeff,
    get_strain_state_dict,
)
from pymatgen.core.elasticity.strain import Deformation, Strain
from pymatgen.core.elasticity.stress import Stress
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
from pymatgen.core.tensors import Tensor
from pymatgen.core.units import FloatWithUnit
from pymatgen.util.testing import TEST_FILES_DIR, MatSciTest

TEST_DIR = f"{TEST_FILES_DIR}/core/elasticity"


class TestElasticTensor(MatSciTest):
    def setup_method(self):
        self.voigt_1 = [
            [59.33, 28.08, 28.08, 0, 0, 0],
            [28.08, 59.31, 28.07, 0, 0, 0],
            [28.08, 28.07, 59.32, 0, 0, 0],
            [0, 0, 0, 26.35, 0, 0],
            [0, 0, 0, 0, 26.35, 0],
            [0, 0, 0, 0, 0, 26.35],
        ]
        mat = np.random.default_rng().standard_normal((6, 6))
        mat += np.transpose(mat)
        self.rand_elastic_tensor = ElasticTensor.from_voigt(mat)
        self.ft = np.array(
            [
                [
                    [[59.33, 0, 0], [0, 28.08, 0], [0, 0, 28.08]],
                    [[0, 26.35, 0], [26.35, 0, 0], [0, 0, 0]],
                    [[0, 0, 26.35], [0, 0, 0], [26.35, 0, 0]],
                ],
                [
                    [[0, 26.35, 0], [26.35, 0, 0], [0, 0, 0]],
                    [[28.08, 0, 0], [0, 59.31, 0], [0, 0, 28.07]],
                    [[0, 0, 0], [0, 0, 26.35], [0, 26.35, 0]],
                ],
                [
                    [[0, 0, 26.35], [0, 0, 0], [26.35, 0, 0]],
                    [[0, 0, 0], [0, 0, 26.35], [0, 26.35, 0]],
                    [[28.08, 0, 0], [0, 28.07, 0], [0, 0, 59.32]],
                ],
            ]
        )

        self.elastic_tensor_1 = ElasticTensor(self.ft)
        filepath = f"{TEST_DIR}/Sn_def_stress.json"
        with open(filepath, "rb") as file:
            self.def_stress_dict = orjson.loads(file.read())
        with open(f"{TEST_DIR}/test_toec_data.json", "rb") as file:
            self.toec_dict = orjson.loads(file.read())
        self.structure = self.get_structure("Sn")

        warnings.simplefilter("always")

    def test_properties(self):
        # compliance tensor
        ct = ComplianceTensor.from_voigt(np.linalg.inv(self.elastic_tensor_1.voigt))
        assert_allclose(ct, self.elastic_tensor_1.compliance_tensor)
        # KG average properties
        assert self.elastic_tensor_1.k_voigt == approx(38.49111111111)
        assert self.elastic_tensor_1.g_voigt == approx(22.05866666666)
        assert self.elastic_tensor_1.k_reuss == approx(38.49110945133)
        assert self.elastic_tensor_1.g_reuss == approx(20.67146635306)
        assert self.elastic_tensor_1.k_vrh == approx(38.49111028122)
        assert self.elastic_tensor_1.g_vrh == approx(21.36506650986)

        # universal anisotropy
        assert self.elastic_tensor_1.universal_anisotropy == approx(0.33553509658699)
        # homogeneous Poisson
        assert self.elastic_tensor_1.homogeneous_poisson == approx(0.26579965576472)
        # voigt notation tensor
        assert_allclose(self.elastic_tensor_1.voigt, self.voigt_1)
        # young's modulus
        assert self.elastic_tensor_1.y_mod == approx(54087787667.160583)

        # prop dict
        prop_dict = self.elastic_tensor_1.property_dict
        assert prop_dict["homogeneous_poisson"] == approx(0.26579965576)
        for k, v in prop_dict.items():
            assert getattr(self.elastic_tensor_1, k) == approx(v)

    def test_directional_elastic_mod(self):
        assert self.elastic_tensor_1.directional_elastic_mod([1, 0, 0]) == approx(self.elastic_tensor_1.voigt[0, 0])
        assert self.elastic_tensor_1.directional_elastic_mod([1, 1, 1]) == approx(73.624444444)

    def test_compliance_tensor(self):
        stress = self.elastic_tensor_1.calculate_stress([0.01] + [0] * 5)
        comp = self.elastic_tensor_1.compliance_tensor
        strain = Strain(comp.einsum_sequence([stress]))
        assert_allclose(strain.voigt, [0.01] + [0] * 5, atol=1e-9)

    def test_directional_poisson_ratio(self):
        v_12 = self.elastic_tensor_1.directional_poisson_ratio([1, 0, 0], [0, 1, 0])
        assert v_12 == approx(0.321388)

    def test_structure_based_methods(self):
        # trans_velocity
        struct = self.structure
        assert self.elastic_tensor_1.trans_v(struct) == approx(1996.3267137182859)
        # long_velocity
        assert self.elastic_tensor_1.long_v(struct) == approx(3534.639656349216)
        # Snyder properties
        assert self.elastic_tensor_1.snyder_ac(struct) == approx(18.061058277029453)
        assert self.elastic_tensor_1.snyder_opt(struct) == approx(0.18937242778808966)
        assert self.elastic_tensor_1.snyder_total(struct) == approx(18.250430704817543)
        # Clarke
        assert self.elastic_tensor_1.clarke_thermalcond(struct) == approx(0.3450388871045757)
        # Cahill
        cahill_thermal_cond = self.elastic_tensor_1.cahill_thermalcond(struct)
        assert cahill_thermal_cond == approx(0.3787576327727649)
        # Agne
        agne_thermal_cond = self.elastic_tensor_1.agne_diffusive_thermalcond(struct)
        assert agne_thermal_cond == approx(0.2380953165347791)
        # Test Agne / Cahill factor
        assert agne_thermal_cond / cahill_thermal_cond == approx(0.6286218307780587)
        # Debye
        assert self.elastic_tensor_1.debye_temperature(struct) == approx(198.80194641857014)

        # structure-property dict
        struct_prop_dict = self.elastic_tensor_1.get_structure_property_dict(struct)
        assert struct_prop_dict["long_v"] == approx(3534.639656349216)
        for val in struct_prop_dict.values():
            assert not isinstance(val, FloatWithUnit)
        for key, val in struct_prop_dict.items():
            if key == "structure":
                assert val == struct
            else:
                attr = getattr(self.elastic_tensor_1, key)
                if callable(attr):
                    assert getattr(self.elastic_tensor_1, key)(struct) == approx(val)
                else:
                    assert getattr(self.elastic_tensor_1, key) == approx(val)

        # Test other sprop dict modes
        struct_prop_dict = self.elastic_tensor_1.get_structure_property_dict(struct, include_base_props=False)
        assert "k_vrh" not in struct_prop_dict

        # Test ValueError being raised for structure properties
        test_et = deepcopy(self.elastic_tensor_1)
        test_et[0][0][0][0] = -100000
        prop_dict = test_et.property_dict
        for attr_name in struct_prop_dict:
            if attr_name not in ([*prop_dict, "structure"]):
                with pytest.raises(
                    ValueError,
                    match="Bulk or shear modulus is negative, property cannot be determined",
                ):
                    getattr(test_et, attr_name)(struct)
        with pytest.raises(
            ValueError,
            match="Bulk or shear modulus is negative, property cannot be determined",
        ):
            test_et.get_structure_property_dict(struct)
        noval_sprop_dict = test_et.get_structure_property_dict(struct, ignore_errors=True)
        assert noval_sprop_dict["snyder_ac"] is None

    def test_new(self):
        assert_allclose(self.elastic_tensor_1, ElasticTensor(self.ft))
        non_symm = self.ft
        non_symm[0, 1, 2, 2] += 1.0
        with pytest.warns(
            UserWarning,
            match="Input elastic tensor does not satisfy standard Voigt symmetries",
        ) as warns:
            ElasticTensor(non_symm)
        assert (
            sum("Input elastic tensor does not satisfy standard Voigt symmetries" in str(warn) for warn in warns) == 1
        )

        bad_tensor1 = np.zeros((3, 3, 3))
        bad_tensor2 = np.zeros((3, 3, 3, 2))
        with pytest.raises(ValueError, match="ElasticTensor input must be rank 4"):
            ElasticTensor(bad_tensor1)
        with pytest.raises(
            ValueError,
            match="Pymatgen only supports 3-dimensional tensors, and default tensor constructor uses standard notation",
        ):
            ElasticTensor(bad_tensor2)

    def test_from_pseudoinverse(self):
        strain_list = [Strain.from_deformation(def_matrix) for def_matrix in self.def_stress_dict["deformations"]]
        stresses = list(self.def_stress_dict["stresses"])
        with pytest.warns(
            UserWarning,
            match="Pseudo-inverse fitting of Strain/Stress lists may yield questionable results from "
            "vasp data, use with caution",
        ):
            et_fl = -0.1 * ElasticTensor.from_pseudoinverse(strain_list, stresses).voigt
        assert_allclose(
            et_fl.round(2),
            [
                [59.29, 24.36, 22.46, 0, 0, 0],
                [28.06, 56.91, 22.46, 0, 0, 0],
                [28.06, 25.98, 54.67, 0, 0, 0],
                [0, 0, 0, 26.35, 0, 0],
                [0, 0, 0, 0, 26.35, 0],
                [0, 0, 0, 0, 0, 26.35],
            ],
        )

    def test_from_independent_strains(self):
        strains = self.toec_dict["strains"]
        stresses = self.toec_dict["stresses"]
        with pytest.warns(UserWarning, match="No eq state found, returning zero voigt stress") as warns:
            et = ElasticTensor.from_independent_strains(strains, stresses)
        assert sum("No eq state found" in str(warn) for warn in warns) == 1
        assert sum("Extra strain states in strain-" in str(warn) for warn in warns) == 1
        assert_allclose(et.voigt, self.toec_dict["C2_raw"], atol=1e1)

    def test_energy_density(self):
        film_elac = ElasticTensor.from_voigt(
            [
                [324.32, 187.3, 170.92, 0, 0, 0],
                [187.3, 324.32, 170.92, 0, 0, 0],
                [170.92, 170.92, 408.41, 0, 0, 0],
                [0, 0, 0, 150.73, 0, 0],
                [0, 0, 0, 0, 150.73, 0],
                [0, 0, 0, 0, 0, 238.74],
            ]
        )

        dfm = Deformation(
            [
                [-9.86004855e-01, 2.27539582e-01, -4.64426035e-17],
                [-2.47802121e-01, -9.91208483e-01, -7.58675185e-17],
                [-6.12323400e-17, -6.12323400e-17, 1.00000000e00],
            ]
        )

        assert film_elac.energy_density(dfm.green_lagrange_strain) == approx(0.00125664672793)

        film_elac.energy_density(
            Strain.from_deformation(
                [
                    [0.99774738, 0.11520994, -0],
                    [-0.11520994, 0.99774738, 0],
                    [-0, -0, 1],
                ]
            )
        )


class TestElasticTensorExpansion(MatSciTest):
    def setup_method(self):
        with open(f"{TEST_DIR}/test_toec_data.json", "rb") as file:
            self.data_dict = orjson.loads(file.read())
        self.strains = [Strain(sm) for sm in self.data_dict["strains"]]
        self.pk_stresses = [Stress(d) for d in self.data_dict["pk_stresses"]]
        self.c2 = self.data_dict["C2_raw"]
        self.c3 = self.data_dict["C3_raw"]
        self.exp = ElasticTensorExpansion.from_voigt([self.c2, self.c3])
        self.cu = Structure.from_spacegroup("Fm-3m", Lattice.cubic(3.623), ["Cu"], [[0] * 3])
        indices = [(0, 0), (0, 1), (3, 3)]
        values = [167.8, 113.5, 74.5]
        cu_c2 = ElasticTensor.from_values_indices(values, indices, structure=self.cu, populate=True)
        indices = [(0, 0, 0), (0, 0, 1), (0, 1, 2), (0, 3, 3), (0, 5, 5), (3, 4, 5)]
        values = [-1507, -965, -71, -7, -901, 45]
        cu_c3 = Tensor.from_values_indices(values, indices, structure=self.cu, populate=True)
        self.exp_cu = ElasticTensorExpansion([cu_c2, cu_c3])
        cu_c4 = Tensor.from_voigt(self.data_dict["Cu_fourth_order"])
        self.exp_cu_4 = ElasticTensorExpansion([cu_c2, cu_c3, cu_c4])

    def test_init(self):
        c_ijkl = Tensor.from_voigt(self.c2)
        c_ijklmn = Tensor.from_voigt(self.c3)
        exp = ElasticTensorExpansion([c_ijkl, c_ijklmn])
        ElasticTensorExpansion.from_voigt([self.c2, self.c3])
        assert exp.order == 3

    def test_from_diff_fit(self):
        ElasticTensorExpansion.from_diff_fit(self.strains, self.pk_stresses)

    def test_calculate_stress(self):
        calc_stress = self.exp.calculate_stress(self.strains[0])
        assert_allclose(self.pk_stresses[0], calc_stress, atol=1e-2)

    def test_energy_density(self):
        e_density = self.exp.energy_density(self.strains[0])
        assert e_density == approx(1.36363099e-4)

    def test_gruneisen(self):
        # Get GGT
        ggt = self.exp_cu.get_ggt([1, 0, 0], [0, 1, 0])
        assert_allclose(ggt, np.eye(3) * np.array([4.92080537, 4.2852349, -0.7147651]))
        # Get TGT
        tgt = self.exp_cu.get_tgt()
        assert_allclose(tgt, np.eye(3) * 2.59631832, atol=1e-12)

        # Get heat capacity
        c0 = self.exp_cu.get_heat_capacity(0, self.cu, [1, 0, 0], [0, 1, 0])
        assert c0 == approx(0.0)
        c = self.exp_cu.get_heat_capacity(300, self.cu, [1, 0, 0], [0, 1, 0])
        assert c == approx(8.285799471733913)

        # Get Gruneisen parameter
        gp = self.exp_cu.get_gruneisen_parameter()
        assert gp == approx(2.59631832)
        _ = self.exp_cu.get_gruneisen_parameter(temperature=200, structure=self.cu)

    def test_thermal_expansion_coeff(self):
        # TODO get rid of duplicates
        alpha_dp = self.exp_cu.thermal_expansion_coeff(self.cu, 300, mode="dulong-petit")
        alpha_dp_ground_truth = 6.3471959e-07 * np.ones((3, 3))
        alpha_dp_ground_truth[np.diag_indices(3)] = 2.2875769e-7
        assert_allclose(alpha_dp_ground_truth, alpha_dp, atol=1e-4)

        alpha_debye = self.exp_cu.thermal_expansion_coeff(self.cu, 300, mode="debye")
        alpha_comp = 5.9435148e-7 * np.ones((3, 3))
        alpha_comp[np.diag_indices(3)] = 21.4533472e-06
        assert_allclose(alpha_comp, alpha_debye, atol=1e-6)

    def test_get_compliance_expansion(self):
        ce_exp = self.exp_cu.get_compliance_expansion()
        et_comp = ElasticTensorExpansion(ce_exp)
        strain_orig = Strain.from_voigt([0.01, 0, 0, 0, 0, 0])
        stress = self.exp_cu.calculate_stress(strain_orig)
        strain_revert = et_comp.calculate_stress(stress)
        assert_allclose(strain_orig, strain_revert, atol=1e-4)

    def test_get_effective_ecs(self):
        # Ensure zero strain is same as SOEC
        test_zero = self.exp_cu.get_effective_ecs(np.zeros((3, 3)))
        assert_allclose(test_zero, self.exp_cu[0])
        strain = np.zeros((3, 3))
        strain[0, 0] = 0.02
        test_2percent = self.exp_cu.get_effective_ecs(strain)
        diff = test_2percent - test_zero
        assert_allclose(self.exp_cu[1].einsum_sequence([strain]), diff)

    def test_get_strain_from_stress(self):
        strain = Strain.from_voigt([0.05, 0, 0, 0, 0, 0])
        stress3 = self.exp_cu.calculate_stress(strain)
        strain_revert3 = self.exp_cu.get_strain_from_stress(stress3)
        assert_allclose(strain, strain_revert3, atol=1e-2)
        # fourth order
        stress4 = self.exp_cu_4.calculate_stress(strain)
        strain_revert4 = self.exp_cu_4.get_strain_from_stress(stress4)
        assert_allclose(strain, strain_revert4, atol=1e-2)

    def test_get_yield_stress(self):
        self.exp_cu_4.get_yield_stress([1, 0, 0])


class TestNthOrderElasticTensor(MatSciTest):
    def setup_method(self):
        with open(f"{TEST_DIR}/test_toec_data.json", "rb") as file:
            self.data_dict = orjson.loads(file.read())
        self.strains = [Strain(sm) for sm in self.data_dict["strains"]]
        self.pk_stresses = [Stress(d) for d in self.data_dict["pk_stresses"]]
        self.c2 = NthOrderElasticTensor.from_voigt(self.data_dict["C2_raw"])
        self.c3 = NthOrderElasticTensor.from_voigt(self.data_dict["C3_raw"])

    def test_init(self):
        c2 = NthOrderElasticTensor(self.c2.tolist())
        c3 = NthOrderElasticTensor(self.c3.tolist())
        c4 = NthOrderElasticTensor(np.zeros([3] * 8))
        for n, c in enumerate([c2, c3, c4]):
            assert c.order == n + 2
        with pytest.raises(ValueError, match="ElasticTensor must have even rank"):
            NthOrderElasticTensor(np.zeros([3] * 5))

    def test_from_diff_fit(self):
        c3 = NthOrderElasticTensor.from_diff_fit(
            self.strains,
            self.pk_stresses,
            eq_stress=self.data_dict["eq_stress"],
            order=3,
        )
        assert_allclose(c3.voigt, self.data_dict["C3_raw"], atol=1e-2)

    def test_calculate_stress(self):
        calc_stress = self.c2.calculate_stress(self.strains[0])
        assert_allclose(self.pk_stresses[0], calc_stress, atol=1e-0)
        # Test calculation from voigt strain
        self.c2.calculate_stress(self.strains[0].voigt)

    def test_energy_density(self):
        self.c3.energy_density(self.strains[0])


class TestDiffFit(MatSciTest):
    """Test various functions related to diff fitting."""

    def setup_method(self):
        with open(f"{TEST_DIR}/test_toec_data.json", "rb") as file:
            self.data_dict = orjson.loads(file.read())
        self.strains = [Strain(sm) for sm in self.data_dict["strains"]]
        self.pk_stresses = [Stress(d) for d in self.data_dict["pk_stresses"]]

    def test_get_strain_state_dict(self):
        strain_inds = [(0,), (1,), (2,), (1, 3), (1, 2, 3)]
        vecs = {}
        strain_states = []
        rng = np.random.default_rng()
        for strain_ind in strain_inds:
            ss = np.zeros(6)
            np.put(ss, strain_ind, 1)
            strain_states.append(tuple(ss))
            vec = np.zeros((4, 6))
            rand_values = rng.uniform(0.1, 1, 4)
            for idx in strain_ind:
                vec[:, idx] = rand_values
            vecs[strain_ind] = vec
        all_strains = [Strain.from_voigt(v).zeroed() for vec in vecs.values() for v in vec]
        rng.shuffle(all_strains)
        all_stresses = [Stress.from_voigt(rng.random(6)).zeroed() for _ in all_strains]
        strain_dict = {k.tobytes(): v for k, v in zip(all_strains, all_stresses, strict=True)}
        ss_dict = get_strain_state_dict(all_strains, all_stresses, add_eq=False)
        # Check length of ss_dict
        assert len(strain_inds) == len(ss_dict)
        # Check sets of strain states are correct
        assert set(strain_states) == set(ss_dict)
        for data in ss_dict.values():
            # Check correspondence of strains/stresses
            for strain, stress in zip(data["strains"], data["stresses"], strict=True):
                assert_allclose(
                    Stress.from_voigt(stress),
                    strain_dict[Strain.from_voigt(strain).tobytes()],
                )
        # Add test to ensure zero strain state doesn't cause issue
        strains, stresses = [Strain.from_voigt([-0.01] + [0] * 5)], [Stress(np.eye(3))]
        ss_dict = get_strain_state_dict(strains, stresses)
        assert_allclose(list(ss_dict), [[1, 0, 0, 0, 0, 0]])

    def test_find_eq_stress(self):
        test_strains = deepcopy(self.strains)
        test_stresses = deepcopy(self.pk_stresses)
        no_eq = find_eq_stress(test_strains, test_stresses)
        assert_allclose(no_eq, np.zeros((3, 3)))
        test_strains[3] = Strain.from_voigt(np.zeros(6))
        eq_stress = find_eq_stress(test_strains, test_stresses)
        assert_allclose(test_stresses[3], eq_stress)

    def test_get_diff_coeff(self):
        def central_diff_weights(Np, ndiv):
            """
            Notes:
                This function is taken from SciPy's deprecated implementation of
                central difference weights. The original code can be found at:
                https://github.com/scipy/scipy/blob/ea916c6f7f487bd53e98de08264
                9d542cc6106ed/scipy/_lib/_finite_differences.py

                License: This code is distributed under the BSD 3-Clause license.

                Copyright (c) 2001-2002 Enthought, Inc. 2003-2024, SciPy Developers.
                All rights reserved.

            Return weights for an Np-point central derivative.

            Assumes equally-spaced function points.

            If weights are in the vector w, then
            derivative is w[0] * f(x-ho*dx) + ... + w[-1] * f(x+h0*dx)

            Args:
                Np (int): Number of points for the central derivative.
                ndiv (int, optional): Number of divisions. Default is 1.

            Returns:
                w (ndarray): Weights for an Np-point central derivative.
                    Its size is `Np`.
            """

            ho = Np >> 1
            x = np.arange(-ho, ho + 1.0)
            x = x[:, np.newaxis]
            X = x**0.0
            for k in range(1, Np):
                X = np.hstack([X, x**k])
            return np.prod(np.arange(1, ndiv + 1), axis=0) * linalg.inv(X)[ndiv]

        forward_11 = get_diff_coeff([0, 1], 1)
        forward_13 = get_diff_coeff([0, 1, 2, 3], 1)
        backward_26 = get_diff_coeff(np.arange(-6, 1), 2)
        central_29 = get_diff_coeff(np.arange(-4, 5), 2)
        assert_allclose(forward_11, [-1, 1])
        assert_allclose(forward_13, [-11 / 6, 3, -3 / 2, 1 / 3])
        assert_allclose(
            backward_26,
            [137 / 180, -27 / 5, 33 / 2, -254 / 9, 117 / 4, -87 / 5, 203 / 45],
        )
        assert_allclose(central_29, central_diff_weights(9, 2))

    def test_generate_pseudo(self):
        strain_states = np.eye(6).tolist()
        for order in (2, 3, 4):
            pseudo_inverses, absent_symbols = generate_pseudo(strain_states, order=order)
            assert len(pseudo_inverses) == order - 1
            assert pseudo_inverses[0].shape == (21, 36)
            assert len(absent_symbols) == len(pseudo_inverses)

    def test_fit(self):
        diff_fit(self.strains, self.pk_stresses, self.data_dict["eq_stress"])
        reduced = [
            (e, pk)
            for e, pk in zip(self.strains, self.pk_stresses, strict=True)
            if not (abs(abs(e) - 0.05) < 1e-10).any()
        ]
        # Get reduced dataset
        r_strains, r_pk_stresses = zip(*reduced, strict=True)
        c2 = diff_fit(r_strains, r_pk_stresses, self.data_dict["eq_stress"], order=2)
        c2, c3, _c4 = diff_fit(r_strains, r_pk_stresses, self.data_dict["eq_stress"], order=4)
        c2, c3 = diff_fit(self.strains, self.pk_stresses, self.data_dict["eq_stress"], order=3)
        c2_red, c3_red = diff_fit(r_strains, r_pk_stresses, self.data_dict["eq_stress"], order=3)
        assert_allclose(c2.voigt, self.data_dict["C2_raw"], atol=1e-5)
        assert_allclose(c3.voigt, self.data_dict["C3_raw"], atol=1e-5)
        assert_allclose(c2, c2_red, atol=1e-0)
        assert_allclose(c3, c3_red, atol=1e1)
