module;

#include "mean_field.h"

export module mean_field:mapping.domain_mapper;
export import :mapping.types;
import :mapping.compactification;
import :utils.user;

export namespace mean_field::mapping {
enum class FaceElementSide : uint8_t { element_1, element_2 };

class ElementDisplacementData {
public:
  ElementDisplacementData(
      const mfem::FiniteElement &element, const mfem::Vector &displacement_dofs,
      mfem::Ordering::Type ordering = mfem::Ordering::byNODES);

  [[nodiscard]] const mfem::FiniteElement &GetElement() const noexcept;
  [[nodiscard]] const mfem::DenseMatrix &GetDofMatrix() const noexcept;
  [[nodiscard]] int GetDimension() const noexcept;
  [[nodiscard]] int GetDofCount() const noexcept;
  [[nodiscard]] mfem::Ordering::Type GetOrdering() const noexcept;

private:
  const mfem::FiniteElement *m_element;
  mfem::DenseMatrix m_dof_matrix;
  int m_dimension;
  mfem::Ordering::Type m_ordering;
};

struct CompactificationPointData {
  double coordinate{0.0};
  mfem::Vector coordinate_gradient;
};

[[nodiscard]] ElementDisplacementData
ElementDisplacementDataFromElementVDofs(const mfem::FiniteElement &element,
                                        const mfem::Vector &displacement_dofs);

class ElementCompactificationData {
public:
  ElementCompactificationData(const mfem::FiniteElement &element,
                              const mfem::Vector &dofs);

  [[nodiscard]] const mfem::FiniteElement &GetElement() const noexcept;
  [[nodiscard]] const mfem::Vector &GetDofs() const noexcept;
  [[nodiscard]] int GetDofCount() const noexcept;

private:
  const mfem::FiniteElement *m_element;
  mfem::Vector m_dofs;
};

struct ElementMappingData {
  const ElementDisplacementData &displacement;
  const ElementCompactificationData &compactification;
};

class DomainMapper {
public:
  class Workspace {
  public:
    explicit Workspace(int dimension = 3);

    void SetDimension(int dimension);

    [[nodiscard]] int GetDimension() const noexcept;

  private:
    friend class DomainMapper;

    int m_dimension;

    mfem::Vector m_shape;
    mfem::DenseMatrix m_mesh_dshape;
    mfem::Vector m_field_value;
    mfem::DenseMatrix m_field_jacobian;

    mfem::Vector m_compactification_shape;
    mfem::DenseMatrix m_compactification_dshape;
    CompactificationPointData m_compactification_point;

    mfem::Vector m_reference_normal;
    mfem::Vector m_mapped_normal;
    mfem::DenseMatrix m_full_element_jacobian;

    mfem::Vector m_vector_temp;
    mfem::DenseMatrix m_matrix_temp_1;
    mfem::DenseMatrix m_matrix_temp_2;

    compactification::ExteriorMapResult m_exterior_result;
    compactification::ExteriorMapVariation m_exterior_variation;
  };

public:
  DomainMapper(
      utils::DomainMapperOptions options,
      std::unique_ptr<const compactification::ExteriorDomainMap> exterior_map);

  DomainMapper(const DomainMapper &) = delete;
  DomainMapper &operator=(const DomainMapper &) = delete;
  DomainMapper(DomainMapper &&) = default;
  DomainMapper &operator=(DomainMapper &&) = default;

  [[nodiscard]] MappingStatus
  EvaluatePoint(const ElementMappingData &element_data,
                mfem::ElementTransformation &transformation,
                const mfem::IntegrationPoint &integration_point,
                Workspace &workspace, MappingPointContext &context) const;

  [[nodiscard]] MappingStatus
  EvaluateVolume(const ElementMappingData &element_data,
                 mfem::ElementTransformation &transformation,
                 const mfem::IntegrationPoint &integration_point,
                 Workspace &workspace, VolumeMappingContext &context) const;

  [[nodiscard]] MappingStatus
  EvaluateFace(const ElementMappingData &element_data,
               mfem::FaceElementTransformations &transformation,
               FaceElementSide side,
               const mfem::IntegrationPoint &integration_point,
               Workspace &workspace, FaceMappingContext &context) const;

  [[nodiscard]] MappingStatus
  EvaluatePointVariation(const ElementMappingData &element_data,
                         const ElementDisplacementData &direction,
                         mfem::ElementTransformation &transformation,
                         const mfem::IntegrationPoint &integration_point,
                         const MappingPointContext &base_context,
                         Workspace &workspace,
                         MappingPointVariation &variation) const;

  [[nodiscard]] MappingStatus
  EvaluateVolumeVariation(const ElementMappingData &element_data,
                          const ElementDisplacementData &direction,
                          mfem::ElementTransformation &transformation,
                          const mfem::IntegrationPoint &integration_point,
                          const VolumeMappingContext &base_context,
                          Workspace &workspace,
                          VolumeMappingVariation &variation) const;

  [[nodiscard]] MappingStatus EvaluateFaceVariation(
      const ElementMappingData &element_data,
      const ElementDisplacementData &direction,
      mfem::FaceElementTransformations &transformation, FaceElementSide side,
      const mfem::IntegrationPoint &integration_point,
      const FaceMappingContext &base_context, Workspace &workspace,
      FaceMappingVariation &variation) const;

  [[nodiscard]] bool IsCompactifiedElement(
      const mfem::ElementTransformation &transformation) const noexcept;
  [[nodiscard]] int GetDimension() const noexcept;
  [[nodiscard]] const compactification::ExteriorDomainMap &
  GetExteriorMap() const noexcept;

private:
  void ValidateElementData(const ElementMappingData &element_data) const;

  void EvaluateField(const ElementDisplacementData &field,
                     mfem::ElementTransformation &transformation,
                     const mfem::IntegrationPoint &integration_point,
                     Workspace &workspace, mfem::Vector &value,
                     mfem::DenseMatrix &jacobian) const;

  [[nodiscard]] MappingStatus EvaluateCompactificationCoordinate(
      const ElementCompactificationData &compactification,
      mfem::ElementTransformation &transformation,
      const mfem::IntegrationPoint &integration_point, Workspace &workspace,
      CompactificationPointData &point_data) const;

  [[nodiscard]] static mfem::ElementTransformation &
  SelectFaceElementTransformation(
      mfem::FaceElementTransformations &transformation, FaceElementSide side);

  [[nodiscard]] static const mfem::IntegrationPoint &
  SelectFaceElementIntegrationPoint(
      mfem::FaceElementTransformations &transformation, FaceElementSide side);

  utils::DomainMapperOptions m_options;
  std::unique_ptr<const compactification::ExteriorDomainMap> m_exterior_map;
};

class GridFunctionMappingEvaluator {
public:
  /*
   * The evaluator references the supplied grid functions and caches copies of
   * their element-local DOFs. Call InvalidateCache() or Refresh() after either
   * grid function's values are modified. Finite-element-space sequence changes
   * are detected automatically.
   *
   * This object owns mutable workspace and cache state and is not thread-safe.
   */
  GridFunctionMappingEvaluator(
      const DomainMapper &mapper,
      const mfem::GridFunction &displacement,
      const mfem::GridFunction &compactification_coordinate);

  /*
   * Discard all element-local field data. The next evaluation reloads its
   * requested element lazily. This operation is idempotent.
   */
  void InvalidateCache() noexcept;

  /*
   * Reload the currently cached element immediately. If no element has been
   * evaluated yet, Refresh() is a validated no-op. If either finite-element
   * space changed sequence, the old element ID is discarded and the next
   * evaluation reloads lazily against the updated spaces.
   */
  void Refresh();

  [[nodiscard]] MappingStatus
  EvaluatePoint(mfem::ElementTransformation &transformation,
                const mfem::IntegrationPoint &integration_point,
                MappingPointContext &context);

  [[nodiscard]] MappingStatus
  EvaluateVolume(mfem::ElementTransformation &transformation,
                 const mfem::IntegrationPoint &integration_point,
                 VolumeMappingContext &context);

  [[nodiscard]] MappingStatus
  EvaluateFace(mfem::FaceElementTransformations &transformation,
               FaceElementSide side,
               const mfem::IntegrationPoint &integration_point,
               FaceMappingContext &context);

  [[nodiscard]] VolumeQuadratureContext
  GetQuadratureContext(mfem::ElementTransformation &transformation,
                       const mfem::IntegrationPoint &integration_point);

  [[nodiscard]] FaceQuadratureContext
  GetFaceQuadratureContext(
      mfem::FaceElementTransformations &transformation,
      const mfem::IntegrationPoint &integration_point,
      FaceElementSide side = FaceElementSide::element_1);

  void GetPhysicalPoint(mfem::ElementTransformation &transformation,
                        const mfem::IntegrationPoint &integration_point,
                        mfem::Vector &physical_position);

private:
  void ValidateFieldBindings() const;
  [[nodiscard]] bool InvalidateForChangedSpaces();
  void LoadElement(int element_id);

  const DomainMapper &m_mapper;
  const mfem::GridFunction &m_displacement;
  const mfem::GridFunction &m_compactification_coordinate;
  const mfem::FiniteElementSpace *m_displacement_space;
  const mfem::FiniteElementSpace *m_compactification_space;
  long m_displacement_space_sequence;
  long m_compactification_space_sequence;
  DomainMapper::Workspace m_workspace;

  mfem::Array<int> m_displacement_dofs;
  mfem::Array<int> m_compactification_dofs;
  mfem::Vector m_element_displacement;
  mfem::Vector m_element_compactification;
  std::unique_ptr<ElementDisplacementData> m_displacement_data;
  std::unique_ptr<ElementCompactificationData> m_compactification_data;
  int m_cached_element_id{-1};
};

} // namespace mean_field::mapping