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

export namespace mean_field::integrators {
    class GravityForceIntegrator : public mfem::BlockNonlinearFormIntegrator {
    public:
        GravityForceIntegrator(const mapping::DomainMapper& map, const mfem::GridFunction& phi);

        void SetPotential(const mfem::GridFunction& phi);

        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:
        const mapping::DomainMapper& m_map;
        const mfem::GridFunction* m_phi;
    };
}