# 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/>.
"""
Tensor mathematics.

All routines operate on numpy.ndarrays of shape (...,3,3).
"""

from warnings import warn as _warn

import numpy as _np
from numpy import typing as _npt


def deviatoric(T: _npt.NDArray[_np.floating]) -> _npt.NDArray[_np.floating]:
    r"""
    Calculate deviatoric part of a tensor.

    Deprecation warning
    -------------------
    .. deprecated:: 3.1.0
        `deviatoric` will be removed in DAMASK 4.0 and replaced by
        `traceless`, `mechanics.deviatoric`, and `mechanics.isochoric`.

    Parameters
    ----------
    T : numpy.ndarray, shape (...,3,3)
        Tensor of which the deviatoric part is computed.

    Returns
    -------
    T' : numpy.ndarray, shape (...,3,3)
        Deviatoric part of T.

    See Also
    --------
    spherical : Calculate spherical part of a tensor.

    Notes
    -----
    The deviatoric part of a tensor is defined as:

    .. math::

        \vb{T}' = \vb{T} - \vb{I}_\text{p},

    where :math:`\vb{I}_\text{p}` is the spherical
    part of the tensor mapped onto identity.
    """
    _warn('tensor.deviatoric will be removed in DAMASK 4.0',
          DeprecationWarning, stacklevel=2)
    return traceless(T)


def eigenvalues(T_sym: _npt.NDArray[_np.floating]) -> _npt.NDArray[_np.floating]:
    r"""
    Calculate eigenvalues of a symmetric tensor.

    Parameters
    ----------
    T_sym : numpy.ndarray, shape (...,3,3)
        Symmetric tensor of which the eigenvalues are computed.

    Returns
    -------
    lambda : numpy.ndarray, shape (...,3)
        Eigenvalues of T_sym sorted in ascending order, each repeated
        according to its multiplicity.

    See Also
    --------
    eigenvectors : Calculate eigenvectors of a symmetric tensor.

    Notes
    -----
    The eigenvalues :math:`\lambda` are defined from:

    .. math::

        \vb{T}_\text{sym} \vb{v}= \lambda \vb{v},

    where :math:`\vb{v}` are the eigenvectors.
    """
    return _np.linalg.eigvalsh(symmetric(T_sym))


def eigenvectors(T_sym: _npt.NDArray[_np.floating],
                 RHS: bool = False) -> _npt.NDArray[_np.floating]:
    r"""
    Calculate eigenvectors of a symmetric tensor.

    Parameters
    ----------
    T_sym : numpy.ndarray, shape (...,3,3)
        Symmetric tensor of which the eigenvectors are computed.
    RHS : bool, optional
        Enforce right-handed coordinate system. Defaults to False.

    Returns
    -------
    v : numpy.ndarray, shape (...,3,3)
        Eigenvectors of T_sym sorted in ascending order of their
        associated eigenvalues.

    See Also
    --------
    eigenvalues : Calculate eigenvalues of a symmetric tensor.

    Notes
    -----
    The eigenvectors :math:`\vb{v}` are defined from:

    .. math::

        \vb{T}_\text{sym} \vb{v}= \lambda \vb{v},

    where :math:`\lambda` are the eigenvalues.
    """
    _,v = _np.linalg.eigh(symmetric(T_sym))

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


def spherical(T: _npt.NDArray[_np.floating],
              tensor: bool = True) -> _npt.NDArray[_np.floating]:
    r"""
    Calculate spherical part of a tensor.

    Parameters
    ----------
    T : numpy.ndarray, shape (...,3,3)
        Tensor of which the spherical part is computed.
    tensor : bool, optional
        Map spherical part onto identity tensor. Defaults to True.

    Returns
    -------
    I_p/p : numpy.ndarray, shape (...,3,3) / shape (...,)
        Spherical part of tensor T. I_p is an isotropic tensor.

    See Also
    --------
    deviatoric : Calculate deviatoric part of a tensor.

    Notes
    -----
    The spherical part of a tensor is defined as:

    .. math::

        p &= \text{trace}( \vb{T} )/3 \\
        \vb{I}_p &= \vb{I} \, p
    """
    sph = _np.trace(T,axis2=-2,axis1=-1)/3.0
    return _np.einsum('...jk,...',_np.eye(3),sph) if tensor else sph


def symmetric(T: _npt.NDArray[_np.floating]) -> _npt.NDArray[_np.floating]:
    r"""
    Symmetrize tensor.

    Parameters
    ----------
    T : numpy.ndarray, shape (...,3,3)
        Tensor of which the symmetrized values are computed.

    Returns
    -------
    T_sym : numpy.ndarray, shape (...,3,3)
        Symmetrized tensor T.

    Notes
    -----
    The symmetrization of a second-order tensor is defined as:

    .. math::

        \vb{T}_\text{sym} = \left( \vb{T}+ \vb{T}^\text{T} \right)/2
    """
    return (T+transpose(T))*0.5


def traceless(T: _npt.NDArray[_np.floating]) -> _npt.NDArray[_np.floating]:
    r"""
    Remove the trace from a tensor.

    Parameters
    ----------
    T : numpy.ndarray, shape (...,3,3)
        Tensor of which the trace is removed.

    Returns
    -------
    T' : numpy.ndarray, shape (...,3,3)
        Tensor T without its trace.

    See Also
    --------
    spherical : Calculate spherical part of a tensor.
    mechanics.deviatoric : Calculate deviatoric part of a
        stress tensor.
    mechanics.isochoric : Calculate isochoric part of a
        strain tensor.

    Notes
    -----
    A stress tensor without trace is called deviatoric;
    a strain tensor without trace is called isochoric.
    """
    return T - spherical(T,tensor=True)


def transpose(T: _npt.NDArray[_np.floating]) -> _npt.NDArray[_np.floating]:
    r"""
    Transpose tensor.

    Parameters
    ----------
    T : numpy.ndarray, shape (...,3,3)
        Tensor of which the transpose is computed.

    Returns
    -------
    T.T : numpy.ndarray, shape (...,3,3)
        Transpose of tensor T.

    Notes
    -----
    The transpose of a second-order tensor is defined as:

    .. math::

        T_{ij}^\text{T} = T_{ji}
    """
    return _np.swapaxes(T,axis2=-2,axis1=-1)
