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

/**
 * @file vanishingstress.hh
 * @brief Defines the VanishingStress material model and related functions.
 * @ingroup  materials
 */

// SPDX-License-Identifier: LGPL-3.0-or-later

#pragma once

#include <ikarus/finiteelements/mechanics/materials/interface.hh>
#include <ikarus/solver/nonlinearsolver/newtonraphson.hh>
#include <ikarus/utils/nonlinearoperator.hh>

namespace Ikarus {

namespace Impl {

  /**
   * @brief Represents a pair of stress matrix indices (row and column).
   */
  struct StressIndexPair
  {
    Eigen::Index row; ///< Row index.
    Eigen::Index col; ///< Column index.
  };

  /**
   * @brief Helper function to create an array of free Voigt indices.
   * @tparam size The size of the fixed pairs array.
   * @param fixed An array of StressIndexPair representing fixed indices.
   * @return std::array<size_t, 6 - size> The array of free Voigt indices.
   */
  template <size_t size>
  consteval auto createfreeVoigtIndices(const std::array<StressIndexPair, size>& fixed) {
    std::array<size_t, 6 - size> res{};
    std::array<size_t, size> voigtFixedIndices;
    std::ranges::transform(fixed, voigtFixedIndices.begin(), [](auto pair) { return toVoigt(pair.row, pair.col); });
    std::ranges::sort(voigtFixedIndices);
    std::ranges::set_difference(std::ranges::iota_view(size_t(0), size_t(6)), voigtFixedIndices, res.begin());
    std::ranges::sort(res);
    return res;
  }

  /**
   * @brief Helper function to create an array of fixed Voigt indices.
   * @tparam size The size of the fixed pairs array.
   * @param fixed An array of StressIndexPair representing fixed indices.
   * @return std::array<size_t, size> The array of fixed Voigt indices.
   */
  template <size_t size>
  consteval auto createFixedVoigtIndices(const std::array<StressIndexPair, size>& fixed) {
    std::array<size_t, size> fixedIndices;
    std::ranges::transform(fixed, fixedIndices.begin(), [](auto pair) { return toVoigt(pair.row, pair.col); });
    std::ranges::sort(fixedIndices);
    return fixedIndices;
  }

  /**
   * @brief Helper function to count the number of diagonal indices in the fixed pairs array.
   * @tparam size The size of the fixed pairs array.
   * @param fixed An array of StressIndexPair representing fixed indices.
   * @return constexpr size_t The number of diagonal indices.
   */
  template <size_t size>
  constexpr size_t countDiagonalIndices(const std::array<StressIndexPair, size>& fixed) {
    size_t count = 0;
    for (auto v : fixed) {
      if (v.col == v.row)
        ++count;
    }
    return count;
  }

} // namespace Impl

/**
 * @brief VanishingStress material model that enforces stress components to be zero.
 * @ingroup materials
 * @tparam stressIndexPair An array of StressIndexPair representing fixed stress components.
 * @tparam MaterialImpl The underlying material model.
 */
template <auto stressIndexPair, typename MaterialImpl>
struct VanishingStress : public Material<VanishingStress<stressIndexPair, MaterialImpl>>
{
  /**
   * @brief Constructor for VanishingStress.
   * @param mat The underlying material model.
   * @param p_tol Tolerance for stress reduction.
   */
  explicit VanishingStress(MaterialImpl mat, typename MaterialImpl::ScalarType p_tol = 1e-12)
      : matImpl{mat},
        tol{p_tol} {}

  using Underlying = MaterialImpl; ///< The underlying material type.

  static constexpr auto fixedPairs        = stressIndexPair;                     ///< Array of fixed stress components.
  static constexpr auto freeVoigtIndices  = createfreeVoigtIndices(fixedPairs);  ///< Free Voigt indices.
  static constexpr auto fixedVoigtIndices = createFixedVoigtIndices(fixedPairs); ///< Fixed Voigt indices.
  static constexpr auto fixedDiagonalVoigtIndicesSize =
      countDiagonalIndices(fixedPairs);                                  ///< Number of fixed diagonal indices.
  static constexpr auto freeStrains = freeVoigtIndices.size();           ///< Number of free strains.
  using ScalarType                  = typename MaterialImpl::ScalarType; ///< Scalar type.

  [[nodiscard]] constexpr std::string nameImpl() const noexcept {
    auto matName = matImpl.name() + "_Vanishing(";
    for (auto p : fixedPairs)
      matName += "(" + std::to_string(p.row) + std::to_string(p.col) + ")";
    matName += ")";
    return matName;
  }

  static constexpr auto strainTag          = MaterialImpl::strainTag;          ///< Strain tag.
  static constexpr auto stressTag          = MaterialImpl::stressTag;          ///< Stress tag.
  static constexpr auto tangentModuliTag   = MaterialImpl::tangentModuliTag;   ///< Tangent moduli tag.
  static constexpr bool energyAcceptsVoigt = MaterialImpl::energyAcceptsVoigt; ///< Energy accepts Voigt notation.
  static constexpr bool stressToVoigt      = true;                             ///< Stress to Voigt notation.
  static constexpr bool stressAcceptsVoigt = true;                             ///< Stress accepts Voigt notation.
  static constexpr bool moduliToVoigt      = true;                             ///< Moduli to Voigt notation.
  static constexpr bool moduliAcceptsVoigt = true;                             ///< Moduli accepts Voigt notation.
  static constexpr double derivativeFactor = 1;                                ///< Derivative factor.

  /**
   * @brief Computes the stored energy for the VanishingStress material.
   * @tparam Derived The derived type of the input matrix.
   * @param E The Green-Lagrangian strain.
   * @return ScalarType The stored energy.
   */
  template <typename Derived>
  ScalarType storedEnergyImpl(const Eigen::MatrixBase<Derived>& E) const {
    const auto [nonOp, Esol] = reduceStress(E);
    return matImpl.storedEnergyImpl(Esol);
  }

  /**
   * @brief Computes the stresses for the VanishingStress material.
   * @tparam voigt A boolean indicating whether to return stresses in Voigt notation.
   * @tparam Derived The derived type of the input matrix.
   * @param E The Green-Lagrangian strain.
   * @return StressMatrix The stresses.
   */
  template <bool voigt, typename Derived>
  auto stressesImpl(const Eigen::MatrixBase<Derived>& E) const {
    const auto [nonOp, Esol] = reduceStress(E);
    auto stressesRed         = matImpl.template stresses<MaterialImpl::strainTag, true>(Esol);

    if constexpr (voigt) {
      return removeCol(stressesRed, fixedVoigtIndices);
    } else
      return fromVoigt(stressesRed, false);
  }

  /**
   * @brief Computes the tangent moduli for the VanishingStress material.
   * @tparam voigt A boolean indicating whether to return tangent moduli in Voigt notation.
   * @tparam Derived The derived type of the input matrix.
   * @param E The Green-Lagrangian strain.
   * @return TangentModuli The tangent moduli.
   */
  template <bool voigt, typename Derived>
  auto tangentModuliImpl(const Eigen::MatrixBase<Derived>& E) const {
    const auto [nonOp, Esol] = reduceStress(E);
    auto C                   = matImpl.template tangentModuli<MaterialImpl::strainTag, true>(Esol);
    if constexpr (voigt)
      return staticCondensation(C, fixedVoigtIndices);
    else
      return fromVoigt(C);
  }

  /**
   * @brief Rebinds the material to a different scalar type.
   * @tparam ScalarTypeOther The target scalar type.
   * @return VanishingStress The rebound VanishingStress material.
   */
  template <typename ScalarTypeOther>
  auto rebind() const {
    auto reboundMatImpl = matImpl.template rebind<ScalarTypeOther>();
    return VanishingStress<stressIndexPair, decltype(reboundMatImpl)>(reboundMatImpl, tol);
  }

private:
  /**
   * @brief Converts the input strain matrix to the appropriate form for stress reduction.
   * @tparam Derived The derived type of the input matrix.
   * @param E The input strain matrix.
   * @return decltype(auto) The converted strain matrix.
   */
  template <typename Derived>
  decltype(auto) maybeFromVoigt(const Eigen::MatrixBase<Derived>& E) const {
    if constexpr (Concepts::EigenVector<Derived>) { // receiving vector means Voigt notation
      return fromVoigt(E.derived(), true);
    } else
      return E.derived();
  }

  /**
   * @brief Initializes unknown strains based on fixed indices.
   * @tparam Derived The derived type of the input matrix.
   * @param E The input strain matrix.
   */
  template <typename Derived>
  void initUnknownStrains(Eigen::MatrixBase<Derived>& E) const {
    for (size_t i = 0; i < fixedPairs.size(); ++i) {
      ScalarType initialVal = E(fixedPairs[i].row, fixedPairs[i].col);
      if constexpr (strainTag == StrainTags::deformationGradient or strainTag == StrainTags::rightCauchyGreenTensor) {
        if (Dune::FloatCmp::eq(initialVal, ScalarType(0.0)) and (fixedPairs[i].row == fixedPairs[i].col))
          initialVal = ScalarType(1.0);
      }
      if (fixedPairs[i].row != fixedPairs[i].col)
        initialVal = ScalarType(0.0);
      E(fixedPairs[i].row, fixedPairs[i].col) = E(fixedPairs[i].col, fixedPairs[i].row) = initialVal;
    }
  }

  /**
   * @brief Reduces stress components to satisfy the vanishing stress condition.
   * @tparam Derived The derived type of the input matrix.
   * @param p_Eraw The input strain matrix.
   * @return std::pair<NonLinearOperator, decltype(auto)> The stress reduction result.
   */
  template <typename Derived>
  auto reduceStress(const Eigen::MatrixBase<Derived>& p_Eraw) const {
    auto E = maybeFromVoigt(p_Eraw);
    initUnknownStrains(E);

    std::array<size_t, fixedDiagonalVoigtIndicesSize> fixedDiagonalVoigtIndices;
    for (size_t ri = 0; auto i : fixedVoigtIndices) {
      auto indexPair = fromVoigt(i);
      if (indexPair[0] == indexPair[1])
        fixedDiagonalVoigtIndices[ri++] = i;
    }

    auto f = [&](auto&) {
      auto S = matImpl.template stresses<MaterialImpl::strainTag, true>(E);
      return S(fixedDiagonalVoigtIndices).eval();
    };
    auto df = [&](auto&) {
      auto moduli = (matImpl.template tangentModuli<MaterialImpl::strainTag, true>(E)).eval();
      return (moduli(fixedDiagonalVoigtIndices, fixedDiagonalVoigtIndices) / MaterialImpl::derivativeFactor).eval();
    };

    auto Er    = E(fixedDiagonalVoigtIndices, fixedDiagonalVoigtIndices).eval().template cast<ScalarType>();
    auto nonOp = Ikarus::NonLinearOperator(functions(f, df), parameter(Er));
    auto nr    = Ikarus::makeNewtonRaphson(
        nonOp, [&](auto& r, auto& A) { return (A.inverse() * r).eval(); },
        [&](auto& /* Ex33 */, auto& Ecomps) {
          for (int ri = 0; auto i : fixedDiagonalVoigtIndices) {
            auto indexPair = fromVoigt(i);
            E(indexPair[0], indexPair[1]) += Ecomps(ri++);
          }
        });
    nr->setup({.tol = tol, .maxIter = 100});
    if (!static_cast<bool>(nr->solve()))
      DUNE_THROW(Dune::MathError, "The stress reduction of material " << nameImpl() << " was unsuccessful\n"
                                                                      << "The strains are\n"
                                                                      << E << "\n The stresses are\n"
                                                                      << f(Er));
    return std::make_pair(nonOp, E);
  }

  MaterialImpl matImpl; ///< The underlying material model.
  double tol{};         ///< Tolerance for stress reduction.
};

/**
 * @brief Factory function to create a VanishingStress material with specified stress indices.
 * @tparam stressIndexPair The array of StressIndexPair representing fixed stress components.
 * @tparam MaterialImpl The underlying material model.
 * @param mat The underlying material model.
 * @param p_tol Tolerance for stress reduction.
 * @return VanishingStress The created VanishingStress material.
 */
template <Impl::StressIndexPair... stressIndexPair, typename MaterialImpl>
auto makeVanishingStress(MaterialImpl mat, typename MaterialImpl::ScalarType p_tol = 1e-12) {
  return VanishingStress<std::to_array({stressIndexPair...}), MaterialImpl>(mat, p_tol);
}

/**
 * @brief Factory function to create a VanishingStress material for plane stress conditions.
 * @tparam MaterialImpl The underlying material model.
 * @param mat The underlying material model.
 * @param p_tol Tolerance for stress reduction.
 * @return VanishingStress The created VanishingStress material for plane stress.
 */
template <typename MaterialImpl>
auto planeStress(const MaterialImpl& mat, typename MaterialImpl::ScalarType p_tol = 1e-8) {
  return makeVanishingStress<Impl::StressIndexPair{2, 1}, Impl::StressIndexPair{2, 0}, Impl::StressIndexPair{2, 2}>(
      mat, p_tol);
}

/**
 * @brief Factory function to create a VanishingStress material for a shell material with zero normal stress
 * condition.
 * @tparam MaterialImpl The underlying material model.
 * @param mat The underlying material model.
 * @param p_tol Tolerance for stress reduction.
 * @return VanishingStress The created VanishingStress material for plane stress.
 */
template <typename MaterialImpl>
auto shellMaterial(const MaterialImpl& mat, typename MaterialImpl::ScalarType p_tol = 1e-8) {
  return makeVanishingStress<Impl::StressIndexPair{2, 2}>(mat, p_tol);
}

/**
 * @brief Factory function to create a VanishingStress material for a beam material with two zero normal stress
 * condition.
 * @tparam MaterialImpl The underlying material model.
 * @param mat The underlying material model.
 * @param p_tol Tolerance for stress reduction.
 * @return VanishingStress The created VanishingStress material for plane stress.
 */
template <typename MaterialImpl>
auto beamMaterial(const MaterialImpl& mat, typename MaterialImpl::ScalarType p_tol = 1e-8) {
  return makeVanishingStress<Impl::StressIndexPair{1, 1}, Impl::StressIndexPair{2, 2}>(mat, p_tol);
}
} // namespace Ikarus
