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

export namespace mean_field::integrators {
    class ContinuityVolumeIntegrator
        : public mfem::BlockNonlinearFormIntegrator {
    public:
        explicit ContinuityVolumeIntegrator(const mapping::DomainMapper &map);

        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;
    };

    class ContinuityFaceIntegrator : public mfem::BlockNonlinearFormIntegrator {
    public:
        explicit ContinuityFaceIntegrator(const mapping::DomainMapper &map);

        void AssembleFaceVector(
            const mfem::Array<const mfem::FiniteElement *> &el1,
            const mfem::Array<const mfem::FiniteElement *> &el2,
            mfem::FaceElementTransformations &Tr,
            const mfem::Array<const mfem::Vector *> &elfun,
            const mfem::Array<mfem::Vector *> &elvect
        ) override;

        void AssembleFaceGrad(
            const mfem::Array<const mfem::FiniteElement *> &el1,
            const mfem::Array<const mfem::FiniteElement *> &el2,
            mfem::FaceElementTransformations &Tr,
            const mfem::Array<const mfem::Vector *> &elfun,
            const mfem::Array2D<mfem::DenseMatrix *> &elmats
        ) override;

    private:
        static bool skip_face(const mfem::FaceElementTransformations &Tr);

        static double compute_u_n(
            const mfem::Vector &v_dofs,
            const mfem::Vector &shape_v_minus,
            const mfem::Vector &n_unit,
            int dof_v_minus,
            int dim
        );

    private:
        const mapping::DomainMapper &m_map;
    };

} // namespace mean_field::integrators