module;
#include "profile.h"
#include <cmath>
#include <cstdint>
#include <limits>
#include <mfem.hpp>

module mean_field;
import :operators.gravity_field;
import :solver.fields;
import :operators.kernels.gravity_field;

namespace {
    using namespace mean_field;
    int get_state_width(const mfem::Array<int> &state_true_offsets) {
        MFEM_VERIFY(
            state_true_offsets.Size() >= 2,
            "The coupled state requires at least one block."
        );
        MFEM_VERIFY(
            state_true_offsets[0] == 0,
            "The coupled state offsets must begin at zero."
        );

        for (int i = 0; i < state_true_offsets.Size() - 1; ++i) {
            MFEM_VERIFY(
                state_true_offsets[i + 1] >= state_true_offsets[i],
                "The coupled state offsets must be nondecreasing."
            );
        }

        MFEM_VERIFY(
            state_true_offsets.Last() > 0, "The coupled state cannot be empty."
        );
        return state_true_offsets.Last();
    }

    int get_gravity_residual_height(const fem::FEM &f) {
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr,
            "GravityFieldOperator requires the gravity-gradient finite-element "
            "space (RT: Raviart-Thomas)."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr,
            "GravityFieldOperator requires the gravity-potential "
            "finite-element "
            "space (L2: Lebesgue "
            "space of square-integrable functions)."
        );
        return f.gravityFluxFes->GetTrueVSize() +
               f.gravityPotentialFes->GetTrueVSize();
    }

    mfem::Array<int> make_gravity_residual_offsets(const fem::FEM &f) {
        mfem::Array<int> offsets(operators::gravity_residual_block_count + 1);
        offsets[0] = 0;
        offsets[1] = f.gravityFluxFes->GetTrueVSize();
        offsets[2] = offsets[1] + f.gravityPotentialFes->GetTrueVSize();
        return offsets;
    }

    template <int index>
    int get_state_block_size(
        const mfem::Array<int> &state_true_offsets,
        const utils::blocks::value_block<index>
    ) {
        MFEM_VERIFY(
            index + 1 < state_true_offsets.Size(),
            "Value block is not present in the state offsets."
        );
        return state_true_offsets[index + 1] - state_true_offsets[index];
    }

    void validate_state_offsets(
        const fem::FEM &f,
        const mfem::Array<int> &state_true_offsets
    ) {
        MFEM_VERIFY(
            f.densityFes != nullptr,
            "GravityFieldOperator requires the density finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldOperator requires the "
                                          "displacement finite-element space."
        );

        using form                   = utils::blocks::gravity_field_form;

        constexpr auto density_block = utils::blocks::get_value_block<form>(
            utils::blocks::density_field.mass_term
        );
        constexpr auto displacement_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::displacement_field.geometry_term
            );
        constexpr auto gravity_gradient_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.gradient_term
            );
        constexpr auto gravity_potential_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.poisson_term
            );

        MFEM_VERIFY(
            state_true_offsets.Size() == form::value_block_count + 1,
            "The gravity state offsets do not match gravity_field_form."
        );
        MFEM_VERIFY(
            get_state_block_size(state_true_offsets, density_block) ==
                f.densityFes->GetTrueVSize(),
            "The density block does not match the density finite-element space."
        );
        MFEM_VERIFY(
            get_state_block_size(state_true_offsets, displacement_block) ==
                f.displacementFes->GetTrueVSize(),
            "The displacement block does not match the displacement "
            "finite-element "
            "space."
        );
        MFEM_VERIFY(
            get_state_block_size(state_true_offsets, gravity_gradient_block) ==
                f.gravityFluxFes->GetTrueVSize(),
            "The gravity-gradient block does not match the RT finite-element "
            "space."
        );
        MFEM_VERIFY(
            get_state_block_size(state_true_offsets, gravity_potential_block) ==
                f.gravityPotentialFes->GetTrueVSize(),
            "The gravity-potential block does not match the potential "
            "finite-element space."
        );
    }

    void validate_gravity_context(const fem::FEM &f) {
        MFEM_VERIFY(
            f.gravityContext.b_form != nullptr,
            "GravityFieldOperator requires the divergence operator."
        );
        MFEM_VERIFY(
            f.gravityContext.BT != nullptr,
            "GravityFieldOperator requires the transpose divergence operator."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr,
            "GravityFieldOperator requires the quadrature-rule factory."
        );
    }

    template <int index>
    mfem::Vector make_read_only_value_view(
        const mfem::Vector &vector,
        const mfem::Array<int> &offsets,
        const utils::blocks::value_block<index>
    ) {
        MFEM_VERIFY(
            index + 1 < offsets.Size(),
            "Value block is not present in the supplied offset array."
        );

        const int begin = offsets[index];
        const int size  = offsets[index + 1] - begin;

        MFEM_VERIFY(
            vector.Size() == offsets.Last(),
            "Vector size does not match the value-block offsets."
        );
        return mfem::Vector(
            const_cast<mfem::real_t *>(vector.GetData()) + begin, size
        );
    }

    template <int index>
    mfem::Vector make_read_only_residual_view(
        const mfem::Vector &vector,
        const mfem::Array<int> &offsets,
        const utils::blocks::residual_block<index> block
    ) {
        const int block_id = block;
        const int begin    = offsets[block_id];
        const int size     = offsets[block_id + 1] - begin;

        MFEM_VERIFY(
            vector.Size() == offsets.Last(),
            "The vector does not match the residual-block layout."
        );

        mfem::Vector view;
        view.MakeRef(const_cast<mfem::Vector &>(vector), begin, size);
        return view;
    }

    template <int index>
    mfem::Vector make_residual_view(
        mfem::Vector &vector,
        const mfem::Array<int> &offsets,
        const utils::blocks::residual_block<index>
    ) {
        MFEM_VERIFY(
            index + 1 < offsets.Size(),
            "Residual block is not present in the supplied offset array."
        );

        const int begin = offsets[index];
        const int size  = offsets[index + 1] - begin;

        MFEM_VERIFY(
            vector.Size() == offsets.Last(),
            "Vector size does not match the residual-block offsets."
        );
        return mfem::Vector(vector.GetData() + begin, size);
    }
} // namespace

namespace mean_field::operators {
    GravityFieldOperator::GravityFieldOperator(
        fem::FEM &f,
        const mapping::DomainMapperStateless &domain_mapper,
        context::gravity_field::GravityFieldLinearizationContext
            &linearization_context,
        const mfem::Array<int> &state_true_offsets,
        GravityFieldJacobianOperator &jacobian
    )
        : Operator(
              get_gravity_residual_height(f),
              get_state_width(state_true_offsets)
          ),
          m_fem(f),
          m_domain_mapper(domain_mapper),
          m_linearization_context(linearization_context),
          m_state_true_offsets(state_true_offsets),
          m_residual_true_offsets(make_gravity_residual_offsets(f)),
          m_jacobian(jacobian) {
        MFEM_VERIFY(f.mesh != nullptr, "GravityFieldOperator requires a mesh.");
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldOperator requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.smesh.exterior_coordinate != nullptr,
            "GravityFieldOperator requires the STROID exterior coordinate."
        );
        MFEM_VERIFY(
            f.smesh.exterior_coordinate->space != nullptr,
            "GravityFieldOperator requires the exterior-coordinate "
            "finite-element "
            "space."
        );
        MFEM_VERIFY(
            f.smesh.exterior_coordinate->values != nullptr,
            "GravityFieldOperator requires the exterior-coordinate values."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "GravityFieldOperator received a domain mapper with the wrong "
            "dimension."
        );

        validate_state_offsets(f, m_state_true_offsets);
        validate_gravity_context(f);

        bool has_vacuum_domain = false;

        for (int i = 0; i < f.mesh->attributes.Size(); ++i) {
            if (f.mesh->attributes[i] ==
                domain_mapper.GetVacuumElementAttribute()) {
                has_vacuum_domain = true;
                break;
            }
        }

        MFEM_VERIFY(
            has_vacuum_domain,
            "GravityFieldOperator requires a compactified vacuum domain."
        );
        MFEM_VERIFY(
            m_residual_true_offsets.Last() == Height(),
            "The gravity residual offsets do not match the operator height."
        );
        MFEM_VERIFY(
            m_state_true_offsets.Last() == Width(),
            "The coupled state offsets do not match the operator width."
        );
    }

    context::gravity_field::GravityFieldPreparationReport
    GravityFieldOperator::Prepare(
        const mfem::Vector &state,
        const context::gravity_field::GravityFieldRevisions &revisions
    ) {
        using form                   = utils::blocks::gravity_field_form;

        constexpr auto density_block = utils::blocks::get_value_block<form>(
            utils::blocks::density_field.mass_term
        );
        constexpr auto displacement_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::displacement_field.geometry_term
            );
        constexpr auto gravity_gradient_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.gradient_term
            );
        constexpr auto gravity_potential_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.poisson_term
            );

        MFEM_VERIFY(
            state.Size() == Width(), "GravityFieldOperator received a "
                                     "preparation state with the wrong size."
        );

        const mfem::Vector density = make_read_only_value_view(
            state, m_state_true_offsets, density_block
        );
        const mfem::Vector displacement = make_read_only_value_view(
            state, m_state_true_offsets, displacement_block
        );
        const mfem::Vector gravity_gradient = make_read_only_value_view(
            state, m_state_true_offsets, gravity_gradient_block
        );
        const mfem::Vector gravity_potential = make_read_only_value_view(
            state, m_state_true_offsets, gravity_potential_block
        );

        return m_linearization_context.Prepare(
            {.density           = density,
             .displacement      = displacement,
             .gravity_gradient  = gravity_gradient,
             .gravity_potential = gravity_potential},
            revisions
        );
    }

    const mfem::Array<int> &
    GravityFieldOperator::GetStateTrueOffsets() const noexcept {
        return m_state_true_offsets;
    }

    const mfem::Array<int> &
    GravityFieldOperator::GetResidualTrueOffsets() const noexcept {
        return m_residual_true_offsets;
    }

    void GravityFieldOperator::ApplyGravityUnknowns(
        const mfem::Vector &gravity_gradient,
        const mfem::Vector &gravity_potential,
        const context::gravity_field::GravityFieldGeometryContext
            &geometry_context,
        mfem::Vector &action
    ) const {
        using form = utils::blocks::gravity_field_form;

        constexpr auto gravity_gradient_residual_block =
            utils::blocks::get_residual_block<form>(
                utils::blocks::gravity_field.gradient_term
            );
        constexpr auto gravity_poisson_residual_block =
            utils::blocks::get_residual_block<form>(
                utils::blocks::gravity_field.poisson_term
            );

        MFEM_VERIFY(
            geometry_context.IsPrepared(),
            "GravityFieldOperator received an unprepared geometry context."
        );
        MFEM_VERIFY(
            gravity_gradient.Size() == m_fem.gravityFluxFes->GetTrueVSize(),
            "GravityFieldOperator received a gravity-gradient vector with the "
            "wrong size."
        );
        MFEM_VERIFY(
            gravity_potential.Size() ==
                m_fem.gravityPotentialFes->GetTrueVSize(),
            "GravityFieldOperator received a gravity-potential vector with the "
            "wrong size."
        );

        action.SetSize(Height());
        action                               = 0.0;

        mfem::Vector gravity_gradient_action = make_residual_view(
            action, m_residual_true_offsets, gravity_gradient_residual_block
        );
        mfem::Vector gravity_poisson_action = make_residual_view(
            action, m_residual_true_offsets, gravity_poisson_residual_block
        );
        mfem::Vector transpose_divergence_action(
            gravity_gradient_action.Size()
        );

        geometry_context.GetMassOperator().Mult(
            gravity_gradient, gravity_gradient_action
        );
        m_fem.gravityContext.BT->Mult(
            gravity_potential, transpose_divergence_action
        );
        gravity_gradient_action += transpose_divergence_action;
        m_fem.gravityContext.b_form->Mult(
            gravity_gradient, gravity_poisson_action
        );
    }

    void GravityFieldOperator::ApplyDensitySource(
        const mfem::Vector &density,
        const context::gravity_field::GravityFieldGeometryContext
            &geometry_context,
        mfem::Vector &action
    ) const {
        using form = utils::blocks::gravity_field_form;

        constexpr auto gravity_poisson_residual_block =
            utils::blocks::get_residual_block<form>(
                utils::blocks::gravity_field.poisson_term
            );

        MFEM_VERIFY(
            geometry_context.IsPrepared(),
            "GravityFieldOperator received an unprepared geometry context."
        );
        MFEM_VERIFY(
            density.Size() == m_fem.densityFes->GetTrueVSize(),
            "GravityFieldOperator received a density vector with the wrong "
            "size."
        );

        action.SetSize(Height());
        action                              = 0.0;

        mfem::Vector gravity_poisson_action = make_residual_view(
            action, m_residual_true_offsets, gravity_poisson_residual_block
        );
        geometry_context.GetSourceOperator().Mult(
            density, gravity_poisson_action
        );
    }

    void GravityFieldOperator::Mult(
        const mfem::Vector &state,
        mfem::Vector &residual
    ) const {
        MEAN_FIELD_PROFILE_SCOPE("GravityFieldOperator::Mult");

        using form                   = utils::blocks::gravity_field_form;

        constexpr auto density_block = utils::blocks::get_value_block<form>(
            utils::blocks::density_field.mass_term
        );
        constexpr auto gravity_gradient_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.gradient_term
            );
        constexpr auto gravity_potential_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.poisson_term
            );

        MFEM_VERIFY(
            state.Size() == Width(),
            "GravityFieldOperator received a state with the wrong size."
        );
        MFEM_VERIFY(
            m_linearization_context.IsPrepared(),
            "GravityFieldOperator must be prepared before Mult is called."
        );

        const mfem::Vector density = make_read_only_value_view(
            state, m_state_true_offsets, density_block
        );
        const mfem::Vector gravity_gradient = make_read_only_value_view(
            state, m_state_true_offsets, gravity_gradient_block
        );
        const mfem::Vector gravity_potential = make_read_only_value_view(
            state, m_state_true_offsets, gravity_potential_block
        );
        const context::gravity_field::GravityFieldGeometryContext
            &geometry_context = m_linearization_context.GetGeometryContext();

        mfem::Vector source;

        ApplyGravityUnknowns(
            gravity_gradient, gravity_potential, geometry_context, residual
        );
        ApplyDensitySource(density, geometry_context, source);

        residual -= source;
    }

    context::gravity_field::GravityFieldLinearizationContext &
    GravityFieldOperator::GetLinearizationContext() noexcept {
        return m_linearization_context;
    }

    const context::gravity_field::GravityFieldLinearizationContext &
    GravityFieldOperator::GetLinearizationContext() const noexcept {
        return m_linearization_context;
    }

    mfem::Operator &
    GravityFieldOperator::GetGradient(const mfem::Vector &state) const {
        MFEM_VERIFY(
            state.Size() == Width(), "GravityFieldOperator received a "
                                     "linearization state with the wrong size."
        );
        MFEM_VERIFY(
            m_linearization_context.IsPrepared(),
            "GravityFieldOperator must be prepared before GetGradient is "
            "called."
        );

        return m_jacobian;
    }
    ReducedGravityFieldOperator::ReducedGravityFieldOperator(
        GravityFieldOperator &gravity_field_operator,
        context::gravity_field::GravityFieldGeometryContext
            &gravity_field_geometry_context,
        const mfem::Vector &displacement
    )
        : Operator(
              gravity_field_operator.Height(),
              gravity_field_operator.Height()
          ),
          m_gravity_field_operator(gravity_field_operator),
          m_gravity_true_offsets(
              gravity_field_operator.GetResidualTrueOffsets()
          ),
          m_gravity_field_geometry_context(gravity_field_geometry_context) {
        using form = utils::blocks::gravity_field_form;

        constexpr auto gravity_gradient_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.gradient_term
            );
        constexpr auto gravity_potential_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::gravity_field.poisson_term
            );
        constexpr auto gravity_gradient_residual_block =
            utils::blocks::get_residual_block<form>(
                utils::blocks::gravity_field.gradient_term
            );
        constexpr auto gravity_poisson_residual_block =
            utils::blocks::get_residual_block<form>(
                utils::blocks::gravity_field.poisson_term
            );

        const mfem::Array<int> &state_offsets =
            m_gravity_field_operator.GetStateTrueOffsets();

        MFEM_VERIFY(
            state_offsets.Size() == form::value_block_count + 1,
            "ReducedGravityFieldOperator received an invalid full-state layout."
        );
        MFEM_VERIFY(
            m_gravity_true_offsets.Size() == form::residual_block_count + 1,
            "ReducedGravityFieldOperator received an invalid gravity-residual "
            "layout."
        );
        MFEM_VERIFY(
            state_offsets[0] == 0, "The full-state offsets must begin at zero."
        );
        MFEM_VERIFY(
            m_gravity_true_offsets[0] == 0,
            "The reduced gravity offsets must begin at zero."
        );
        MFEM_VERIFY(
            state_offsets.Last() == m_gravity_field_operator.Width(),
            "The full-state offsets do not match the gravity-field operator "
            "width."
        );
        MFEM_VERIFY(
            m_gravity_true_offsets.Last() == m_gravity_field_operator.Height(),
            "The reduced gravity offsets do not match the gravity-field "
            "operator "
            "height."
        );
        MFEM_VERIFY(
            Width() == Height(), "ReducedGravityFieldOperator must be square."
        );

        const int full_gradient_size =
            state_offsets[static_cast<int>(gravity_gradient_block) + 1] -
            state_offsets[gravity_gradient_block];
        const int full_potential_size =
            state_offsets[static_cast<int>(gravity_potential_block) + 1] -
            state_offsets[gravity_potential_block];
        const int reduced_gradient_size =
            m_gravity_true_offsets
                [static_cast<int>(gravity_gradient_residual_block) + 1] -
            m_gravity_true_offsets[gravity_gradient_residual_block];
        const int reduced_potential_size =
            m_gravity_true_offsets
                [static_cast<int>(gravity_poisson_residual_block) + 1] -
            m_gravity_true_offsets[gravity_poisson_residual_block];

        MFEM_VERIFY(
            full_gradient_size == reduced_gradient_size,
            "The reduced gravity-gradient block does not match the full-state "
            "gravity-gradient block."
        );
        MFEM_VERIFY(
            full_potential_size == reduced_potential_size,
            "The reduced gravity-potential block does not match the Poisson "
            "residual block."
        );

        SetDisplacement(displacement);
    }

    void ReducedGravityFieldOperator::SetDisplacement(
        const mfem::Vector &displacement
    ) {
        ValidateDisplacement(displacement);

        context::gravity_field::DiscretizationRevision discretization_revision;
        context::gravity_field::DisplacementRevision displacement_revision;

        if (m_gravity_field_geometry_context.IsPrepared()) {
            discretization_revision =
                m_gravity_field_geometry_context.GetDiscretizationRevision();
            displacement_revision =
                m_gravity_field_geometry_context.GetDisplacementRevision();

            MFEM_VERIFY(
                displacement_revision.value <
                    std::numeric_limits<std::uint64_t>::max(),
                "The reduced gravity displacement revision has overflowed."
            );
            ++displacement_revision.value;
        }

        m_gravity_field_geometry_context.Prepare(
            displacement, discretization_revision, displacement_revision
        );
    }

    const mfem::Vector &ReducedGravityFieldOperator::GetDisplacement() const {
        return m_gravity_field_geometry_context.GetDisplacement();
    }

    void ReducedGravityFieldOperator::BuildRightHandSide(
        const mfem::Vector &density,
        mfem::Vector &right_hand_side
    ) const {
        ValidateDensity(density);

        m_gravity_field_operator.ApplyDensitySource(
            density, m_gravity_field_geometry_context, right_hand_side
        );

        MFEM_VERIFY(
            right_hand_side.Size() == Height(),
            "ReducedGravityFieldOperator produced a right-hand side with the "
            "wrong "
            "size."
        );
    }

    void ReducedGravityFieldOperator::Mult(
        const mfem::Vector &gravity_state,
        mfem::Vector &action
    ) const {
        MEAN_FIELD_PROFILE_SCOPE("ReducedGravityFieldOperator::Mult");

        using form = utils::blocks::gravity_field_form;

        constexpr auto gravity_gradient_residual_block =
            utils::blocks::get_residual_block<form>(
                utils::blocks::gravity_field.gradient_term
            );
        constexpr auto gravity_poisson_residual_block =
            utils::blocks::get_residual_block<form>(
                utils::blocks::gravity_field.poisson_term
            );

        ValidateGravityState(gravity_state);

        const mfem::Vector gravity_gradient_true = make_read_only_residual_view(
            gravity_state, m_gravity_true_offsets,
            gravity_gradient_residual_block
        );
        const mfem::Vector gravity_potential_true =
            make_read_only_residual_view(
                gravity_state, m_gravity_true_offsets,
                gravity_poisson_residual_block
            );

        m_gravity_field_operator.ApplyGravityUnknowns(
            gravity_gradient_true, gravity_potential_true,
            m_gravity_field_geometry_context, action
        );

        MFEM_VERIFY(
            action.Size() == Height(), "ReducedGravityFieldOperator produced "
                                       "an action with the wrong size."
        );
    }

    GravityFieldOperator &
    ReducedGravityFieldOperator::GetGravityFieldOperator() noexcept {
        return m_gravity_field_operator;
    }

    const GravityFieldOperator &
    ReducedGravityFieldOperator::GetGravityFieldOperator() const noexcept {
        return m_gravity_field_operator;
    }

    context::gravity_field::GravityFieldGeometryContext &
    ReducedGravityFieldOperator::GetGeometryContext() noexcept {
        return m_gravity_field_geometry_context;
    }

    const context::gravity_field::GravityFieldGeometryContext &
    ReducedGravityFieldOperator::GetGeometryContext() const noexcept {
        return m_gravity_field_geometry_context;
    }

    const mfem::Array<int> &
    ReducedGravityFieldOperator::GetGravityTrueOffsets() const noexcept {
        return m_gravity_true_offsets;
    }

    void ReducedGravityFieldOperator::ValidateDisplacement(
        const mfem::Vector &displacement
    ) const {
        using form = utils::blocks::gravity_field_form;

        constexpr auto displacement_block =
            utils::blocks::get_value_block<form>(
                utils::blocks::displacement_field.geometry_term
            );

        const mfem::Array<int> &state_offsets =
            m_gravity_field_operator.GetStateTrueOffsets();
        const int expected_size =
            state_offsets[static_cast<int>(displacement_block) + 1] -
            state_offsets[displacement_block];

        MFEM_VERIFY(
            displacement.Size() == expected_size,
            "ReducedGravityFieldOperator received a displacement with the "
            "wrong "
            "size."
        );

        for (int i = 0; i < displacement.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(displacement(i)),
                "ReducedGravityFieldOperator received a non-finite "
                "displacement "
                "value."
            );
        }
    }

    void ReducedGravityFieldOperator::ValidateDensity(
        const mfem::Vector &density
    ) const {
        using form                   = utils::blocks::gravity_field_form;

        constexpr auto density_block = utils::blocks::get_value_block<form>(
            utils::blocks::density_field.mass_term
        );

        const mfem::Array<int> &state_offsets =
            m_gravity_field_operator.GetStateTrueOffsets();
        const int expected_size =
            state_offsets[static_cast<int>(density_block) + 1] -
            state_offsets[density_block];

        MFEM_VERIFY(
            density.Size() == expected_size,
            "ReducedGravityFieldOperator received a density with the wrong "
            "size."
        );
    }

    void ReducedGravityFieldOperator::ValidateGravityState(
        const mfem::Vector &gravity_state
    ) const {
        MFEM_VERIFY(
            gravity_state.Size() == Width(),
            "ReducedGravityFieldOperator received "
            "a gravity state with the wrong size."
        );
    }
} // namespace mean_field::operators