module;
#include <cstdint>
#include <mfem.hpp>
export module mean_field:integrators.gravity;
import :mapping.domain_mapper;

export namespace mean_field::integrators {
    enum class GravityForceJacobianMode : std::uint8_t { minimal, field_coupled, exact };

    class GravityMomentumIntegrator : public mfem::BlockNonlinearFormIntegrator {
    public:
        explicit GravityMomentumIntegrator(
            const mapping::DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            GravityForceJacobianMode jacobian_mode = GravityForceJacobianMode::field_coupled
        );

        void SetJacobianMode(GravityForceJacobianMode jacobian_mode);
        void SetIntegrationRule(const mfem::IntegrationRule &integration_rule);

        [[nodiscard]] GravityForceJacobianMode GetJacobianMode() const;

        void AssembleElementVector(
            const mfem::Array<const mfem::FiniteElement *> &el,
            mfem::ElementTransformation &Tr,
            const mfem::Array<const mfem::Vector *> &elfun,
            const mfem::Array<mfem::Vector *> &elvec
        ) override;
        void AssembleElementGrad(
            const mfem::Array<const mfem::FiniteElement *> &el,
            mfem::ElementTransformation &Tr,
            const mfem::Array<const mfem::Vector *> &elfun,
            const mfem::Array2D<mfem::DenseMatrix *> &elmats
        ) override;

    private:
        mapping::GridFunctionMappingEvaluator m_mapping;
        GravityForceJacobianMode m_jacobian_mode;
        const mfem::IntegrationRule *m_integration_rule{nullptr};
    };
} // namespace mean_field::integrators