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

/**
 * @file ferequirements.hh
 * @brief Definition of the LinearElastic class for finite element mechanics computations.
 * @ingroup finiteelements
 */

#pragma once

#include <iosfwd>
#include <map>
#include <set>
#include <vector>

#include <dune/common/exceptions.hh>

#include <Eigen/Core>

#include <ikarus/utils/makeenum.hh>

namespace Ikarus {
  // clang-format off
  /**
 * ScalarAffordances
 * \ingroup Affordancestags
 * \brief A strongly typed enum class representing the scalar affordance
 */
  MAKE_ENUM(ScalarAffordances,
            noAffordance,
            mechanicalPotentialEnergy,
            microMagneticPotentialEnergy
      );

  /**
* VectorAffordances
* \ingroup Affordancestags
* \brief A strongly typed enum class representing the vector affordance
*/
  MAKE_ENUM(VectorAffordances,
            noAffordance,
            forces,
            microMagneticForces
      );

  /**
* MatrixAffordances
* \ingroup Affordancestags
* \brief A strongly typed enum class representing the matrix affordance
*/
  MAKE_ENUM(MatrixAffordances,
            noAffordance,
            stiffness,
            materialstiffness,
            geometricstiffness,
            stiffnessdiffBucklingVector,
            microMagneticHessian,
            mass
      );

  /**
*
* \ingroup FEParameterTags
* \brief A strongly typed enum class representing the FE parameter
*/
  MAKE_ENUM(FEParameter,
            noParameter,
            loadfactor,
            time
      );

  /**
*
* \ingroup FEParameterTags
* \brief A strongly typed enum class representing the type of the solutions vectors
*/
  MAKE_ENUM(FESolutions,
            noSolution,
            displacement,
            velocity,
            director,
            magnetizationAndVectorPotential
      );

  /**
*
* \ingroup FEParameterTags
* \brief A strongly typed enum class representing the type of the result request
*/
  MAKE_ENUM(ResultType,
            noType,
            magnetization,
            gradientNormOfMagnetization,
            vectorPotential,
            divergenceOfVectorPotential,
            BField,
            HField,
            cauchyStress,
            PK2Stress,
            linearStress,
            director
      );

  // clang-format on

  /**
   * \brief Struct representing a collection of affordances.
   */
  struct AffordanceCollectionImpl {
    ScalarAffordances scalarAffordances{ScalarAffordances::noAffordance};
    VectorAffordances vectorAffordances{VectorAffordances::noAffordance};
    MatrixAffordances matrixAffordances{MatrixAffordances::noAffordance};
  };

  /**
   * \brief Concept to check if a given type is one of the predefined affordance enums or the AffordanceCollectionImpl.
   */
  template <typename Type>
  concept FEAffordance
      = std::is_same_v<std::remove_cvref_t<Type>, ScalarAffordances> or std::is_same_v<std::remove_cvref_t<Type>,
                                                                                       VectorAffordances> or std::
          is_same_v<std::remove_cvref_t<Type>, MatrixAffordances> or std::is_same_v<std::remove_cvref_t<Type>,
                                                                                    AffordanceCollectionImpl>;

  inline constexpr VectorAffordances forces = VectorAffordances::forces;

  inline constexpr MatrixAffordances stiffness                   = MatrixAffordances::stiffness;
  inline constexpr MatrixAffordances stiffnessdiffBucklingVector = MatrixAffordances::stiffnessdiffBucklingVector;
  inline constexpr MatrixAffordances mass                        = MatrixAffordances::mass;

  inline constexpr ScalarAffordances potentialEnergy = ScalarAffordances::mechanicalPotentialEnergy;

  namespace AffordanceCollections {
    inline constexpr AffordanceCollectionImpl elastoStatics
        = {ScalarAffordances::mechanicalPotentialEnergy, VectorAffordances::forces, MatrixAffordances::stiffness};
  }

  namespace Impl {
    template <typename T>
    struct DeduceRawVectorType {
      static_assert(!std::is_same<T, T>::value, "You should end up in the provided specializations");
    };

    template <typename T>
    struct DeduceRawVectorType<std::reference_wrapper<T>> {
      using Type = T;
    };

    template <typename T>
    struct DeduceRawVectorType<Eigen::Ref<T>> {
      using Type = Eigen::Ref<T>;
    };
  }  // namespace Impl

  /**
   * \class FErequirements
   * \brief Class representing the requirements for finite element calculations.
   *
   * This class defines the requirements for finite element calculations, including the types of solution vectors
   * and parameters needed. It provides methods to add affordances, insert parameters, and manage global solution
   * vectors.
   *
   * \tparam SolutionVectorType_ Type of the solution vector, defaulting to std::reference_wrapper<Eigen::VectorXd>.
   * \tparam ParameterType_ Type of the parameter, defaulting to std::reference_wrapper<double>.
   *
   */
  template <typename SolutionVectorType_ = std::reference_wrapper<Eigen::VectorXd>,
            typename ParameterType_      = std::reference_wrapper<double>>
  class FErequirements {
  public:
    using SolutionVectorType    = SolutionVectorType_;
    using SolutionVectorTypeRaw = Impl::DeduceRawVectorType<std::remove_cvref_t<SolutionVectorType_>>::Type;
    using ParameterType         = ParameterType_;
    using ParameterTypeRaw      = ParameterType_::type;

    /**
     * \brief Add an affordance to the requirements.
     *
     * This function adds the specified affordance to the requirements.
     *
     * \tparam Affordance Type of affordance to be added.
     * \param affordance The affordance to be added.
     * \return Reference to the updated FErequirements instance.
     */
    template <FEAffordance Affordance>
    FErequirements& addAffordance(Affordance&& affordance) {
      if constexpr (std::is_same_v<Affordance, ScalarAffordances>)
        affordances.scalarAffordances = affordance;
      else if constexpr (std::is_same_v<Affordance, VectorAffordances>)
        affordances.vectorAffordances = affordance;
      else if constexpr (std::is_same_v<Affordance, MatrixAffordances>)
        affordances.matrixAffordances = affordance;
      else if constexpr (std::is_same_v<Affordance, AffordanceCollectionImpl>)
        affordances = affordance;
      return *this;
    }

    /**
     * \brief Insert a parameter into the requirements.
     *
     * This function inserts the specified parameter into the requirements.
     *
     * \param key The key representing the parameter.
     * \param val Reference to the raw parameter value.
     * \return Reference to the updated FErequirements instance.
     */
    FErequirements& insertParameter(const FEParameter& key, ParameterTypeRaw& val) {
      parameter.insert_or_assign(key, val);
      return *this;
    }

    /**
     * \brief Insert a global solution vector into the requirements.
     *
     * This function inserts the specified global solution vector into the requirements.
     *
     * \param key The key representing the type of the solution vector.
     * \param sol Reference to the raw global solution vector.
     * \return Reference to the updated FErequirements instance.
     */
    FErequirements& insertGlobalSolution(const FESolutions& key, SolutionVectorTypeRaw& sol) {
      sols.insert_or_assign(key, sol);
      return *this;
    }

    /**
     * \brief Get the raw global solution vector for a specific type.
     *
     * This function retrieves the raw global solution vector for the specified type.
     *
     * \param key The key representing the type of the solution vector.
     * \return Const reference to the raw global solution vector.
     *
     * \throws Dune::RangeError if the specified type is not found in the requirements.
     */
    const SolutionVectorTypeRaw& getGlobalSolution(const FESolutions& key) const {
      try {
        if constexpr (std::is_same_v<SolutionVectorType, std::reference_wrapper<Eigen::VectorXd>>)
          return sols.at(key).get();
        else
          return sols.at(key);
      } catch (std::out_of_range& oor) {
        DUNE_THROW(Dune::RangeError, std::string("Out of Range error: ") + std::string(oor.what())
                                         + " in getGlobalSolution with key" + toString(key));
        abort();
      }
    }

    /**
     * \brief Get the raw parameter value for a specific key.
     *
     * This function retrieves the raw parameter value for the specified key.
     *
     * \param key The key representing the parameter.
     * \return Const reference to the raw parameter value.
     *
     * \throws Dune::RangeError if the specified key is not found in the requirements.
     */
    const ParameterTypeRaw& getParameter(FEParameter&& key) const {
      try {
        return parameter.at(key).get();
      } catch (std::out_of_range& oor) {
        DUNE_THROW(Dune::RangeError, std::string("Out of Range error: ") + std::string(oor.what())
                                         + " in getParameter with key" + toString(key));
        abort();
      }
    }

    /**
     * \brief Check if a specific affordance is present in the requirements.
     *
     * This function checks if the specified affordance is present in the requirements.
     *
     * \tparam Affordance Type of affordance to be checked.
     * \param affordance The affordance to be checked.
     * \return True if the affordance is present, false otherwise.
     */
    template <FEAffordance Affordance>
    bool hasAffordance(Affordance&& affordance) const {
      if constexpr (std::is_same_v<Affordance, ScalarAffordances>)
        return affordances.scalarAffordances == affordance;
      else if constexpr (std::is_same_v<Affordance, VectorAffordances>)
        return affordances.vectorAffordances == affordance;
      else if constexpr (std::is_same_v<Affordance, MatrixAffordances>)
        return affordances.matrixAffordances == affordance;
      else if constexpr (std::is_same_v<Affordance, AffordanceCollectionImpl>)
        return affordances == affordance;
    }

  private:
    std::map<FESolutions, SolutionVectorType> sols;
    std::map<FEParameter, ParameterType> parameter;
    AffordanceCollectionImpl affordances;
  };

  /**
   * \class ResultTypeMap
   * \brief Class representing a map of result types to result arrays.
   *
   * This class provides a mapping between different result types and their corresponding result arrays.
   * It allows inserting or assigning results and retrieving results based on the result type.
   *
   * \tparam ParameterType Type of the parameters, defaulting to double.
   *
   */
  template <typename ParameterType = double>
  class ResultTypeMap {
  public:
    using ResultArray = Eigen::Matrix<ParameterType, Eigen::Dynamic, Eigen::Dynamic, 0, 9, 3>;

    /**
     * \brief Insert or assign a result to the map.
     *
     * This function inserts or assigns the specified result array to the map with the given result type.
     *
     * \param resultType The type of the result to be inserted or assigned.
     * \param resultArray The result array to be inserted or assigned.
     */
    void insertOrAssignResult(ResultType&& resultType, const ResultArray& resultArray) {
      results.insert_or_assign(resultType, resultArray);
    }

    /**
     * \brief Get the result array for a specific result type.
     *
     * This function retrieves the result array for the specified result type.
     *
     * \param resultType The type of the result to be retrieved.
     * \return Reference to the result array.
     *
     * \throws std::out_of_range if the specified result type is not found in the map.
     */
    ResultArray& getResult(const ResultType& resultType) { return results.at(resultType); }

    /**
     * \brief Get the result array for a single result type.
     *
     * This function retrieves the result array when only a single result type is present in the map.
     *
     * \return Reference to the result array.
     *
     * \throws Dune::RangeError if the map does not contain a single result.
     */
    auto& getSingleResult() {
      if (results.size() != 1)
        DUNE_THROW(Dune::RangeError, "getSingleResult can only be called when a single result was inserted");
      else
        return *(results.begin());
    }

  private:
    std::map<ResultType, ResultArray> results;
  };

  template <typename Type>
  concept ResultTypeConcept = std::is_same_v<Type, ResultType>;

  /**
   * \class ResultRequirements
   * \brief Class representing the requirements for obtaining specific results.
   *
   * This class encapsulates the requirements for obtaining results, including the desired result types,
   * associated affordances, and input parameters. It is templated on the type of FErequirements.
   *
   * \tparam FErequirements Type representing the finite element requirements.
   *                       Default is FErequirements<>.
   *
   */
  template <typename FErequirements = FErequirements<>>
  class ResultRequirements {
  public:
    using ParameterTypeRaw      = typename FErequirements::ParameterTypeRaw;
    using SolutionVectorType    = typename FErequirements::SolutionVectorType;
    using SolutionVectorTypeRaw = typename FErequirements::SolutionVectorTypeRaw;

    /**
     * \brief Constructor with FErequirements and result types.
     *
     * Constructs a ResultRequirements object with the given FErequirements and set of result types.
     *
     * \param req Finite element requirements.
     * \param p_resType Set of result types.
     */
    ResultRequirements(FErequirements&& req, std::set<ResultType>&& p_resType)
        : reqB{req}, resType(std::move(p_resType)) {}

    /**
     * \brief Constructor with only FErequirements.
     *
     * Constructs a ResultRequirements object with the given FErequirements and an empty set of result types.
     *
     * \param req Finite element requirements.
     */
    explicit ResultRequirements(const FErequirements& req) : reqB{req} {}

    /**
     * \brief Default constructor.
     *
     * Constructs an empty ResultRequirements object.
     */
    ResultRequirements() = default;

    /**
     * \brief Check if a specific result type is requested.
     *
     * Checks if the specified result type is requested.
     *
     * \param key The result type to check.
     * \return True if the result type is requested, false otherwise.
     */
    bool isResultRequested(ResultType&& key) const { return resType.contains(key); }

    /**
     * \brief Add an affordance to the finite element requirements.
     *
     * Adds the specified affordance to the finite element requirements.
     *
     * \tparam Affordance Type of affordance to be added.
     * \param affordance The affordance to be added.
     * \return Reference to the ResultRequirements object.
     */
    template <FEAffordance Affordance>
    ResultRequirements& addAffordance(Affordance&& affordance) {
      reqB.addAffordance(std::forward<Affordance>(affordance));
      return *this;
    }

    /**
     * \brief Insert a parameter into the finite element requirements.
     *
     * Inserts the specified parameter into the finite element requirements.
     *
     * \param key The parameter key.
     * \param val The parameter value.
     * \return Reference to the ResultRequirements object.
     */
    ResultRequirements& insertParameter(FEParameter&& key, ParameterTypeRaw& val) {
      reqB.insertParameter(std::forward<FEParameter>(key), val);
      return *this;
    }

    /**
     * \brief Insert a global solution into the finite element requirements.
     *
     * Inserts the specified global solution into the finite element requirements.
     *
     * \param key The global solution key.
     * \param sol The global solution value.
     * \return Reference to the ResultRequirements object.
     */
    ResultRequirements& insertGlobalSolution(FESolutions&& key, SolutionVectorTypeRaw& sol) {
      reqB.insertGlobalSolution(std::forward<FESolutions>(key), sol);
      return *this;
    }

    /**
     * \brief Add one or more result types to the set of requested results.
     *
     * Adds the specified result types to the set of requested results.
     *
     * \tparam ResultTypes Types of results to be added.
     * \param keys The result types to be added.
     * \return Reference to the ResultRequirements object.
     */
    template <ResultTypeConcept... ResultTypes>
    ResultRequirements& addResultRequest(ResultTypes&&... keys) {
      resType.insert({std::move(keys)...});
      return *this;
    }

    /**
     * \brief Get the global solution for a specific global solution type.
     *
     * Retrieves the global solution for the specified global solution type.
     *
     * \param key The global solution type.
     * \return Reference to the global solution value.
     */
    const SolutionVectorTypeRaw& getGlobalSolution(FESolutions&& key) const {
      return reqB.getGlobalSolution(std::move(key));
    }

    /**
     * \brief Get the value of a specific parameter.
     *
     * Retrieves the value of the specified parameter.
     *
     * \param key The parameter key.
     * \return Reference to the parameter value.
     */
    const ParameterTypeRaw& getParameter(FEParameter&& key) const { return reqB.getParameter(std::move(key)); }

    /**
     * \brief Get the associated finite element requirements.
     *
     * Retrieves the associated finite element requirements.
     *
     * \return Reference to the finite element requirements.
     */
    const FErequirements& getFERequirements() const { return reqB; }

    /**
     * \brief Get the requested result type.
     *
     * Retrieves the requested result type when only a single result type is present in the set of requested results.
     *
     * \return Reference to the requested result type.
     *
     * \throws Dune::InvalidStateException if the set of requested results does not contain a single result type.
     */
    auto getRequestedResult() const {
      if (resType.size() == 1)
        return *(resType.begin());
      else {
        DUNE_THROW(Dune::InvalidStateException, "This function can only be called when a single result is requested");
      }
    }

  private:
    std::set<ResultType> resType;
    FErequirements reqB;
  };
}  // namespace Ikarus
