# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright 2024-2025 Max-Planck-Institut für Nachhaltige Materialien GmbH
# Copyright 2011-2024 Max-Planck-Institut für Eisenforschung GmbH
#
# DAMASK is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
import pytest
import numpy as np

from damask import tensor

def traceless(T):
    return T - spherical(T)

def eigenvalues(T_sym):
    return np.linalg.eigvalsh(symmetric(T_sym))

def eigenvectors(T_sym,RHS=False):
    (u,v) = np.linalg.eigh(symmetric(T_sym))

    if RHS:
        if np.linalg.det(v) < 0.0: v[:,2] *= -1.0
    return v

def symmetric(T):
    return (T+transpose(T))*0.5

def transpose(T):
    return T.T

def spherical(T,tensor=True):
    sph = np.trace(T)/3.0
    return sph if not tensor else np.eye(3)*sph


n = 1000

@pytest.mark.parametrize('vectorized,single',[(tensor.traceless,    traceless),
                                              (tensor.eigenvalues,  eigenvalues),
                                              (tensor.eigenvectors, eigenvectors),
                                              (tensor.symmetric,    symmetric),
                                              (tensor.transpose,    transpose),
                                              (tensor.spherical,    spherical),
                                             ])
def test_vectorize_1_arg(np_rng,vectorized,single):
    epsilon     = np_rng.random((n,3,3))
    epsilon_vec = np.reshape(epsilon,(n//10,10,3,3))
    for i,v in enumerate(np.reshape(vectorized(epsilon_vec),vectorized(epsilon).shape)):
        assert np.allclose(single(epsilon[i]),v)

def test_symmetric(np_rng):
    """Ensure that a symmetric tensor is half of the sum of a tensor and its transpose."""
    x = np_rng.random((n,3,3))
    assert np.allclose(tensor.symmetric(x)*2.0,tensor.transpose(x)+x)

def test_transpose(np_rng):
    """Ensure that a symmetric tensor equals its transpose."""
    x = tensor.symmetric(np_rng.random((n,3,3)))
    assert np.allclose(tensor.transpose(x),x)

def test_eigenvalues(np_rng):
    """Ensure that the characteristic polynomial can be solved."""
    A = tensor.symmetric(np_rng.random((n,3,3)))
    lambd = tensor.eigenvalues(A)
    s = np_rng.integers(n)
    for i in range(3):
        assert np.allclose(np.linalg.det(A[s]-lambd[s,i]*np.eye(3)),.0)

def test_eigenvalues_and_vectors(np_rng):
    """Ensure that eigenvalues and -vectors are the solution to the characteristic polynomial."""
    A = tensor.symmetric(np_rng.random((n,3,3)))
    lambd = tensor.eigenvalues(A)
    x     = tensor.eigenvectors(A)
    s = np_rng.integers(n)
    for i in range(3):
        assert np.allclose(np.dot(A[s]-lambd[s,i]*np.eye(3),x[s,:,i]),.0)

def test_eigenvectors_RHS(np_rng):
    """Ensure that RHS coordinate system does only change sign of determinant."""
    A = tensor.symmetric(np_rng.random((n,3,3)))
    LRHS = np.linalg.det(tensor.eigenvectors(A,RHS=False))
    RHS  = np.linalg.det(tensor.eigenvectors(A,RHS=True))
    assert np.allclose(np.abs(LRHS),RHS)

def test_spherical_traceless_part(np_rng):
    """Ensure that full tensor is sum of spherical and traceless part."""
    x = np_rng.random((n,3,3))
    assert np.allclose(tensor.spherical(x,True) + tensor.traceless(x),
                       x)
def test_spherical_mapping(np_rng):
    """Ensure that mapping to tensor is correct."""
    x = np_rng.random((n,3,3))
    tnsr   = tensor.spherical(x,True)
    scalar = tensor.spherical(x,False)
    assert np.allclose(np.linalg.det(tnsr),
                       scalar**3.0)

def test_traceless(np_rng):
    T = np_rng.random((n,3,3))
    assert (np.trace(tensor.traceless(T),axis1=1,axis2=2) < 5.e-16).all()

def test_traceless_diagonal(np_rng):
    I_n = np.broadcast_to(np.eye(3),(n,3,3))
    r   = np.logical_not(I_n)*np_rng.random((n,3,3))
    assert np.allclose(tensor.traceless(I_n+r),r)
