// SPDX-FileCopyrightText: 2021-2024 The Ikarus Developers mueller@ibb.uni-stuttgart.de
// SPDX-License-Identifier: LGPL-3.0-or-later

/**
 * \file tensorutils.hh
 * \brief Helper for the Eigen::Tensor types
 */

#pragma once

#include <numeric>
#include <ranges>
#include <unsupported/Eigen/CXX11/Tensor>

#include <dune/common/promotiontraits.hh>

#include <ikarus/utils/concepts.hh>
#include <ikarus/utils/math.hh>
namespace Ikarus {

/**
 * @brief View an Eigen matrix as an Eigen Tensor with specified dimensions.
 * \ingroup tensor
 * @tparam Derived Type of the input Eigen matrix.
 * @tparam T Type of the elements in the matrix.
 * @tparam rank Rank of the resulting Tensor.
 * @param matrix Input Eigen matrix to be cast.
 * @param dims Dimensions of the resulting Tensor.
 * @return Eigen::Tensor<typename Derived::Scalar, rank> The casted Eigen Tensor.
 */
template <typename Derived, typename T, auto rank>
Eigen::Tensor<typename Derived::Scalar, rank> tensorView(const Eigen::EigenBase<Derived>& matrix,
                                                         const std::array<T, rank>& dims) {
  return Eigen::TensorMap<const Eigen::TensorFixedSize<
      const typename Derived::Scalar, Eigen::Sizes<Derived::RowsAtCompileTime, Derived::ColsAtCompileTime>>>(
      matrix.derived().eval().data(), dims);
}

/**
 * @brief Computes the dyadic product of two Eigen tensors.
 * \details The components of the result read  \f[ \CI_{ijkl} = A_{ij}B_{kl}. \f]
 * \ingroup tensor
 * @param A_ij First tensor.
 * @param B_kl Second tensor.
 * @return  Resulting tensor after the dyadic product.
 */
auto dyadic(const auto& A_ij, const auto& B_kl) {
  Eigen::array<Eigen::IndexPair<long>, 0> empty_index_list = {};
  return A_ij.contract(B_kl, empty_index_list).eval();
}

/**
 * @brief Generates a symmetric identity fourth-order tensor.
 * \ingroup tensor
 * @tparam ScalarType Type of the elements in the tensor.
 * @tparam dim Dimension of the tensor.
 * @return  Symmetric identity fourth-order tensor.
 */
template <typename ScalarType = double, int dim = 3>
auto symmetricIdentityFourthOrder() {
  Eigen::TensorFixedSize<ScalarType, Eigen::Sizes<dim, dim, dim, dim>> idTensor;
  for (int i = 0; i < dim; ++i)
    for (int j = 0; j < dim; ++j)
      for (int k = 0; k < dim; ++k)
        for (int l = 0; l < dim; ++l)
          idTensor(i, j, k, l) = 0.5 * ((i == k) * (j == l) + (i == l) * (j == k));
  return idTensor;
}

/**
 * @brief Generates a symmetric fourth-order tensor based on two second-order input tensors.
 * \ingroup tensor
 * \details The components of the result read  \f[ \CI_{ijkl} = \frac{1}{2} \left(A_{ik}B_{jl}+A_{il}B_{jk} \right) .
 * \f]
 * @tparam ScalarType Type of the elements in the tensors.
 * @tparam dim Dimension of the tensors.
 * @param A First tensor.
 * @param B Second tensor.
 * @return  Symmetric fourth-order tensor.
 */
template <typename ScalarType = double, int dim = 3>
auto symmetricFourthOrder(const auto& A, const auto& B) {
  Eigen::TensorFixedSize<ScalarType, Eigen::Sizes<dim, dim, dim, dim>> res;
  for (int i = 0; i < dim; ++i)
    for (int j = 0; j < dim; ++j)
      for (int k = 0; k < dim; ++k)
        for (int l = 0; l < dim; ++l)
          res(i, j, k, l) = 0.5 * (A(i, k) * B(j, l) + A(i, l) * B(j, k));
  return res;
}

/**
 * @brief Generates an identity fourth-order tensor.
 *  \ingroup tensor
 * \details The components of the result read  \f[ \CI_{ijkl} = \de_{ij}\de_{kl}. \f]
 * @tparam ScalarType Type of the elements in the tensor.
 * @tparam dim Dimension of the tensor.
 * @return  Identity fourth-order tensor.
 */
template <typename ScalarType = double, int dim = 3>
auto identityFourthOrder() {
  Eigen::TensorFixedSize<ScalarType, Eigen::Sizes<dim, dim, dim, dim>> idTensor;
  idTensor.setZero();
  for (int i = 0; i < dim; ++i)
    for (int k = 0; k < dim; ++k)
      idTensor(i, i, k, k) = 1.0;
  return idTensor;
}

/**
 * @brief Computes the IKJL product of two matrices.
 *  \ingroup tensor
 * \details The components of the result read  \f[ \CI_{ijkl} = A_{ik}B_{jl}, \f] which simply swaps the inner slots
 * `j`and `k`
 * @tparam AType Type of the first matrix.
 * @tparam BType Type of the second matrix.
 * @param A First matrix.
 * @param B Second matrix.
 * @return  Resulting tensor of the IKJL product.
 */
template <typename AType, typename BType>
auto fourthOrderIKJL(const Eigen::MatrixBase<AType>& A, const Eigen::MatrixBase<BType>& B) {
  static_assert(AType::RowsAtCompileTime == BType::RowsAtCompileTime);
  static_assert(AType::ColsAtCompileTime == BType::ColsAtCompileTime);
  using ScalarResultType = typename Dune::PromotionTraits<typename AType::Scalar, typename BType::Scalar>::PromotedType;
  constexpr int dim      = AType::RowsAtCompileTime;
  Eigen::TensorFixedSize<ScalarResultType, Eigen::Sizes<dim, dim, dim, dim>> res;
  for (int i = 0; i < dim; ++i)
    for (int j = 0; j < dim; ++j)
      for (int k = 0; k < dim; ++k)
        for (int l = 0; l < dim; ++l)
          res(i, j, k, l) = A(i, k) * B(j, l);
  return res;
}

/**
 * @brief Creates a symmetric fourth-order tensor in the two specified slots of the input tensor.
 *  \ingroup tensor
 * @tparam ScalarType Type of the elements in the tensor.
 * @param t Input tensor.
 * @param slots Indices of the slots to be swapped.
 * @return  Fourth-order Tensor which is symmetric in the given slots.
 */
template <typename ScalarType, long int dim>
auto symTwoSlots(const Eigen::TensorFixedSize<ScalarType, Eigen::Sizes<dim, dim, dim, dim>>& t,
                 const std::array<size_t, 2>& slots) {
  std::array<size_t, 4> shuffleSlot;
  std::iota(shuffleSlot.begin(), shuffleSlot.end(), 0);    // create 0,1,2,3 array
  std::swap(shuffleSlot[slots[0]], shuffleSlot[slots[1]]); // swap  the given slots
  return (0.5 * (t + t.shuffle(shuffleSlot))).eval();
}

/**
 * @brief Converts 2D indices to Voigt notation index.
 *  \ingroup tensor
 * @param i Row index.
 * @param j Column index.
 * @return Eigen::Index Voigt notation index.
 *
 * This function converts 2D indices (i, j) to a Voigt notation index.
 * The Voigt notation is used to represent the six unique components of a symmetric 3x3 matrix
 * in a one-dimensional array.
 *
 * If the input indices are not within the valid range (0, 1, 2), an assertion failure is triggered.
 */
constexpr Eigen::Index toVoigt(Eigen::Index i, Eigen::Index j) noexcept {
  if (i == j) // _00 -> 0, _11 -> 1,  _22 -> 2
    return i;
  if ((i == 1 and j == 2) or (i == 2 and j == 1)) // _12 and _21 --> 3
    return 3;
  if ((i == 0 and j == 2) or (i == 2 and j == 0)) // _02 and _20 --> 4
    return 4;
  if ((i == 0 and j == 1) or (i == 1 and j == 0)) // _01 and _10 --> 5
    return 5;
  assert(i < 3 and j < 3 && "For Voigt notation the indices need to be 0,1 or 2.");
  __builtin_unreachable();
}

/**
 * @brief Converts a fourth-order tensor of fixed size 3x3x3x3 to a Voigt notation matrix of size 6x6.
 *  \ingroup tensor
 * @tparam ScalarType Data type of the tensor elements.
 * @param ft Fourth-order tensor .
 * @return Voigt notation matrix.
 *
 * This function converts a fourth-order tensor to a Voigt notation matrix, which is a symmetric 6x6 matrix
 * containing the unique components of the input tensor. The mapping from the tensor indices to the Voigt notation
 * indices is performed by the toVoigt function.
 *
 * \remarks The current implementation
 * does not take advantage of this symmetry.
 */
template <typename ScalarType = double>
Eigen::Matrix<ScalarType, 6, 6> toVoigt(const Eigen::TensorFixedSize<ScalarType, Eigen::Sizes<3, 3, 3, 3>>& ft) {
  Eigen::Matrix<ScalarType, 6, 6> mat;
  for (Eigen::Index i = 0; i < 3; ++i)
    for (Eigen::Index j = 0; j < 3; ++j)
      for (Eigen::Index k = 0; k < 3; ++k)
        for (Eigen::Index l = 0; l < 3; ++l)
          mat(toVoigt(i, j), toVoigt(k, l)) = ft(i, j, k, l);
  return mat;
}

/**
 * @brief Converts a square 2x2 or 3x3 matrix to a Voigt notation vector.
 *  \ingroup tensor
 * @tparam ST Data type of the matrix elements.
 * @tparam size Number of rows and columns of the square matrix.
 * @tparam Options Eigen matrix options.
 * @param E Input matrix of size (size x size).
 * @param isStrain Flag indicating whether the conversion is for strain (true) or not (false).
 * @return  Vector with components in Voigt notation vector.
 *
 * This function converts a square matrix to a Voigt notation vector, which contains the unique components of
 * the input matrix. The mapping from the matrix indices to the Voigt notation indices is performed by the toVoigt
 * function.
 *
 * The optional isStrain parameter allows the user to specify whether the conversion is intended for strain
 * calculations. If isStrain is true, the off-diagonal components are multiplied by 2, providing the correct Voigt
 * notation for symmetric strain tensors.
 */
template <typename ST, int size, int Options>
requires(size > 0 and size <= 3)
auto toVoigt(const Eigen::Matrix<ST, size, size, Options, size, size>& E, bool isStrain = true) {
  Eigen::Vector<ST, (size * (size + 1)) / 2> EVoigt;
  EVoigt.template segment<size>(0) = E.diagonal();

  const ST possibleStrainFactor = isStrain ? 2.0 : 1.0;
  if constexpr (size == 2)
    EVoigt(2) = E(0, 1) * possibleStrainFactor;
  else if constexpr (size == 3) {
    EVoigt(size)     = E(1, 2) * possibleStrainFactor;
    EVoigt(size + 1) = E(0, 2) * possibleStrainFactor;
    EVoigt(size + 2) = E(0, 1) * possibleStrainFactor;
  }
  return EVoigt;
}

/**
 * @brief Converts a vector given in Voigt notation to a matrix.
 *  \ingroup tensor
 * @tparam ST Scalar type of the vector elements.
 * @tparam size Size of the Voigt notation vector.
 * @param EVoigt Voigt notation vector.
 * @param isStrain Flag indicating whether the vector represents a strain (default is true).
 * @return Matrix corresponding to the vector in Voigt notation.
 *  \details
 * This function converts a vector given in Voigt notation to the corresponding matrix. The conversion depends on the
 * size The parameter `isStrain` is used to determine the conversion factor for off-diagonal components, which need to
 * be divided by 2 in the matrix representation if the quantity is a strain tensor.
 *
 * The function requires that the size of the Voigt notation vector is valid (1, 3, or 6).
 */
template <typename ST, int size>
requires(size == 1 or size == 3 or size == 6)
auto fromVoigt(const Eigen::Vector<ST, size>& EVoigt, bool isStrain = true) {
  constexpr int matrixSize = (-1 + ct_sqrt(1 + 8 * size)) / 2;
  Eigen::Matrix<ST, matrixSize, matrixSize> E;
  E.diagonal() = EVoigt.template segment<3>(0);

  const ST possibleStrainFactor = isStrain ? 0.5 : 1.0;
  if constexpr (matrixSize == 2)
    E(0, 1) = E(1, 0) = EVoigt(2) * possibleStrainFactor;
  else if constexpr (matrixSize == 3) {
    E(2, 1) = E(1, 2) = EVoigt(matrixSize) * possibleStrainFactor;
    E(2, 0) = E(0, 2) = EVoigt(matrixSize + 1) * possibleStrainFactor;
    E(1, 0) = E(0, 1) = EVoigt(matrixSize + 2) * possibleStrainFactor;
  }
  return E;
}

/**
 * @brief Converts a Voigt notation index to matrix indices.
 *  \ingroup tensor
 * @param i Voigt notation index.
 * @return Matrix indices corresponding to the Voigt notation index.
 *  \details
 * This function converts a Voigt notation index to the corresponding matrix indices. The mapping is based on the
 * assumption that the Voigt notation indices 0, 1, and 2 represent the diagonal components `00`, `11`, and `22`,
 * respectively. The remaining Voigt notation indices (3, 4, and 5) correspond to the off-diagonal components
 * (`12` and `21`, `02` and `20`, `01` and `10`).
 *
 * The function asserts that the input index is within the valid range for Voigt notation (0 to 5).
 */
constexpr std::array<size_t, 2> fromVoigt(size_t i) {
  if (i < 3) // _00 -> 0, _11 -> 1,  _22 -> 2
    return {i, i};
  else if (i == 3)
    return {1, 2};
  else if (i == 4)
    return {0, 2};
  else if (i == 5)
    return {0, 1};
  else {
    assert(i < 6 && "For Voigt notation the indices need to be 0 and 5.");
    __builtin_unreachable();
  }
}

/**
 * @brief Converts a matrix in Voigt notation to a Fourth-order tensor.
 *  \ingroup tensor
 * @tparam ScalarType Scalar type of the matrix elements.
 * @param CVoigt Voigt notation matrix.
 * @return Fourth-order tensor corresponding to the matrix in Voigt notation.
 *  \details
 * This function converts a Voigt notation matrix to the corresponding 4th-order tensor. The function uses the
 * `fromVoigt` function to map matrix indices to tensor indices. The resulting tensor is symmetric due to symmetry
 * considerations.
 */
template <typename ScalarType>
auto fromVoigt(const Eigen::Matrix<ScalarType, 6, 6>& CVoigt) {
  Eigen::TensorFixedSize<ScalarType, Eigen::Sizes<3, 3, 3, 3>> C;
  // size_t iR=0,jR=0;
  for (size_t i = 0; i < 6; ++i) {
    for (size_t j = 0; j < 6; ++j) {
      auto firstIndices                                                       = fromVoigt(i);
      auto secondIndices                                                      = fromVoigt(j);
      C(firstIndices[0], firstIndices[1], secondIndices[0], secondIndices[1]) = CVoigt(i, j);
    }
  }
  return C;
}

} // namespace Ikarus
