#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

using namespace mean_field;

TEST_CASE(
    "Gravity Force Integrator Jacobian Matches Residual Linearization",
    tags::gravity_integrator_unit
) {
    constexpr int dim                       = 3;
    constexpr double finite_difference_step = 1.0e-3;
    constexpr double jacobian_tolerance     = 1.0e-8;
    constexpr double zero_tolerance         = 1.0e-14;

    constexpr int velocity_block            = solver::block_index(solver::FieldBlock::velocity);
    constexpr int density_block             = solver::block_index(solver::FieldBlock::density);
    constexpr int gravity_gradient_block    = solver::block_index(solver::FieldBlock::gravity_gradient);
    constexpr int gravity_potential_block   = solver::block_index(solver::FieldBlock::gravity_potential);
    constexpr int displacement_block        = solver::block_index(solver::FieldBlock::displacement);
    constexpr int block_count               = solver::field_block_count;

    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(1, 1, 1, mfem::Element::HEXAHEDRON, 1.0, 1.0, 1.0);

    mfem::H1_FECollection velocity_fec(2, dim);
    mfem::L2_FECollection density_fec(1, dim);
    mfem::RT_FECollection gravity_gradient_fec(1, dim);
    mfem::L2_FECollection gravity_potential_fec(1, dim);
    mfem::H1_FECollection displacement_fec(2, dim);
    mfem::H1_FECollection compactification_fec(1, dim);

    mfem::FiniteElementSpace velocity_fes(&mesh, &velocity_fec, dim, mfem::Ordering::byVDIM);
    mfem::FiniteElementSpace density_fes(&mesh, &density_fec);
    mfem::FiniteElementSpace gravity_gradient_fes(&mesh, &gravity_gradient_fec);
    mfem::FiniteElementSpace gravity_potential_fes(&mesh, &gravity_potential_fec);
    mfem::FiniteElementSpace displacement_fes(&mesh, &displacement_fec, dim, mfem::Ordering::byVDIM);
    mfem::FiniteElementSpace compactification_fes(&mesh, &compactification_fec);

    mfem::GridFunction displacement(&displacement_fes);
    displacement = 0.0;
    mfem::GridFunction compactification_coordinate(&compactification_fes);
    compactification_coordinate                          = 0.0;
    mapping::DomainMapper domain_mapper                  = field_dof_test_utils::make_domain_mapper();

    const mfem::FiniteElement *velocity_element          = velocity_fes.GetFE(0);
    const mfem::FiniteElement *density_element           = density_fes.GetFE(0);
    const mfem::FiniteElement *gravity_gradient_element  = gravity_gradient_fes.GetFE(0);
    const mfem::FiniteElement *gravity_potential_element = gravity_potential_fes.GetFE(0);
    const mfem::FiniteElement *displacement_element      = displacement_fes.GetFE(0);
    mfem::ElementTransformation *transformation          = mesh.GetElementTransformation(0);

    const int velocity_dofs_count                        = velocity_element->GetDof();
    const int density_dofs_count                         = density_element->GetDof();
    const int gravity_gradient_dofs_count                = gravity_gradient_element->GetDof();
    const int gravity_potential_dofs_count               = gravity_potential_element->GetDof();
    const int displacement_dofs_count                    = displacement_element->GetDof();
    const int velocity_size                              = dim * velocity_dofs_count;
    const int displacement_size                          = dim * displacement_dofs_count;

    mfem::Vector velocity_dofs(velocity_size);
    mfem::Vector density_dofs(density_dofs_count);
    mfem::Vector gravity_gradient_dofs(gravity_gradient_dofs_count);
    mfem::Vector gravity_potential_dofs(gravity_potential_dofs_count);
    mfem::Vector displacement_dofs(displacement_size);
    velocity_dofs          = 0.0;
    gravity_potential_dofs = 0.0;
    displacement_dofs      = 0.0;

    for (int i = 0; i < density_dofs_count; ++i) {
        density_dofs(i) = 0.8 + 0.07 * static_cast<double>(i + 1);
    }

    for (int i = 0; i < gravity_gradient_dofs_count; ++i) {
        const double sign        = i % 2 == 0 ? 1.0 : -1.0;
        gravity_gradient_dofs(i) = sign * 0.04 * static_cast<double>(i + 1);
    }

    mfem::Vector density_direction(density_dofs_count);
    mfem::Vector gravity_gradient_direction(gravity_gradient_dofs_count);
    mfem::Vector velocity_direction(velocity_size);
    mfem::Vector displacement_direction(displacement_size);

    for (int i = 0; i < density_dofs_count; ++i) {
        density_direction(i) = 0.13 - 0.02 * static_cast<double>(i);
    }

    for (int i = 0; i < gravity_gradient_dofs_count; ++i) {
        const double sign             = i % 3 == 0 ? -1.0 : 1.0;
        gravity_gradient_direction(i) = sign * (0.03 + 0.005 * static_cast<double>(i));
    }

    for (int i = 0; i < velocity_size; ++i) {
        velocity_direction(i) = 0.01 * static_cast<double>(i + 1);
    }

    for (int i = 0; i < displacement_size; ++i) {
        displacement_direction(i) = -0.008 * static_cast<double>(i + 1);
    }

    mfem::Array<const mfem::FiniteElement *> elements(block_count);
    elements[velocity_block]          = velocity_element;
    elements[density_block]           = density_element;
    elements[gravity_gradient_block]  = gravity_gradient_element;
    elements[gravity_potential_block] = gravity_potential_element;
    elements[displacement_block]      = displacement_element;

    mfem::Array<const mfem::Vector *> element_state(block_count);
    element_state[velocity_block]          = &velocity_dofs;
    element_state[density_block]           = &density_dofs;
    element_state[gravity_gradient_block]  = &gravity_gradient_dofs;
    element_state[gravity_potential_block] = &gravity_potential_dofs;
    element_state[displacement_block]      = &displacement_dofs;

    mfem::Vector velocity_residual(velocity_size);
    mfem::Vector density_residual(density_dofs_count);
    mfem::Vector gravity_gradient_residual(gravity_gradient_dofs_count);
    mfem::Vector gravity_potential_residual(gravity_potential_dofs_count);
    mfem::Vector displacement_residual(displacement_size);

    mfem::Array<mfem::Vector *> element_residual(block_count);
    element_residual[velocity_block]          = &velocity_residual;
    element_residual[density_block]           = &density_residual;
    element_residual[gravity_gradient_block]  = &gravity_gradient_residual;
    element_residual[gravity_potential_block] = &gravity_potential_residual;
    element_residual[displacement_block]      = &displacement_residual;

    integrators::GravityMomentumIntegrator integrator(
        domain_mapper, displacement, compactification_coordinate, integrators::GravityForceJacobianMode::field_coupled
    );

    const int maximum_order = std::max(
        velocity_element->GetOrder(), std::max(density_element->GetOrder(), gravity_gradient_element->GetOrder())
    );
    const mfem::IntegrationRule &integration_rule =
        mfem::IntRules.Get(velocity_element->GetGeomType(), 2 * maximum_order + 8);
    integrator.SetIntegrationRule(integration_rule);

    mfem::DenseMatrix dv_dv(velocity_size, velocity_size);
    mfem::DenseMatrix dv_drho(velocity_size, density_dofs_count);
    mfem::DenseMatrix dv_dgrad_phi(velocity_size, gravity_gradient_dofs_count);
    mfem::DenseMatrix dv_ddisplacement(velocity_size, displacement_size);
    dv_dv            = 1.0;
    dv_drho          = 1.0;
    dv_dgrad_phi     = 1.0;
    dv_ddisplacement = 1.0;

    mfem::Array2D<mfem::DenseMatrix *> element_matrices(block_count, block_count);

    for (int row = 0; row < block_count; ++row) {
        for (int column = 0; column < block_count; ++column) {
            element_matrices(row, column) = nullptr;
        }
    }

    element_matrices(velocity_block, velocity_block)         = &dv_dv;
    element_matrices(velocity_block, density_block)          = &dv_drho;
    element_matrices(velocity_block, gravity_gradient_block) = &dv_dgrad_phi;
    element_matrices(velocity_block, displacement_block)     = &dv_ddisplacement;

    integrator.AssembleElementGrad(elements, *transformation, element_state, element_matrices);

    auto assemble_velocity_residual = [&](const mfem::Vector &density_state,
                                          const mfem::Vector &gravity_gradient_state) {
        element_state[density_block]          = &density_state;
        element_state[gravity_gradient_block] = &gravity_gradient_state;
        integrator.AssembleElementVector(elements, *transformation, element_state, element_residual);
        return mfem::Vector(velocity_residual);
    };

    auto relative_error = [](mfem::Vector computed, const mfem::Vector &reference) {
        computed -= reference;
        return computed.Norml2() / std::max(reference.Norml2(), 1.0e-30);
    };

    mfem::Vector density_plus(density_dofs);
    mfem::Vector density_minus(density_dofs);
    density_plus.Add(finite_difference_step, density_direction);
    density_minus.Add(-finite_difference_step, density_direction);

    mfem::Vector density_residual_plus  = assemble_velocity_residual(density_plus, gravity_gradient_dofs);
    mfem::Vector density_residual_minus = assemble_velocity_residual(density_minus, gravity_gradient_dofs);
    mfem::Vector density_finite_difference(density_residual_plus);
    density_finite_difference -= density_residual_minus;
    density_finite_difference *= 0.5 / finite_difference_step;

    mfem::Vector density_jacobian_action(velocity_size);
    dv_drho.Mult(density_direction, density_jacobian_action);

    mfem::Vector gravity_gradient_plus(gravity_gradient_dofs);
    mfem::Vector gravity_gradient_minus(gravity_gradient_dofs);
    gravity_gradient_plus.Add(finite_difference_step, gravity_gradient_direction);
    gravity_gradient_minus.Add(-finite_difference_step, gravity_gradient_direction);

    mfem::Vector gravity_residual_plus  = assemble_velocity_residual(density_dofs, gravity_gradient_plus);
    mfem::Vector gravity_residual_minus = assemble_velocity_residual(density_dofs, gravity_gradient_minus);
    mfem::Vector gravity_finite_difference(gravity_residual_plus);
    gravity_finite_difference -= gravity_residual_minus;
    gravity_finite_difference *= 0.5 / finite_difference_step;

    mfem::Vector gravity_jacobian_action(velocity_size);
    dv_dgrad_phi.Mult(gravity_gradient_direction, gravity_jacobian_action);

    mfem::Vector combined_density_plus(density_dofs);
    mfem::Vector combined_density_minus(density_dofs);
    mfem::Vector combined_gravity_plus(gravity_gradient_dofs);
    mfem::Vector combined_gravity_minus(gravity_gradient_dofs);
    combined_density_plus.Add(finite_difference_step, density_direction);
    combined_density_minus.Add(-finite_difference_step, density_direction);
    combined_gravity_plus.Add(finite_difference_step, gravity_gradient_direction);
    combined_gravity_minus.Add(-finite_difference_step, gravity_gradient_direction);

    mfem::Vector combined_residual_plus  = assemble_velocity_residual(combined_density_plus, combined_gravity_plus);
    mfem::Vector combined_residual_minus = assemble_velocity_residual(combined_density_minus, combined_gravity_minus);
    mfem::Vector combined_finite_difference(combined_residual_plus);
    combined_finite_difference -= combined_residual_minus;
    combined_finite_difference *= 0.5 / finite_difference_step;

    mfem::Vector combined_jacobian_action(density_jacobian_action);
    combined_jacobian_action += gravity_jacobian_action;

    mfem::Vector inactive_velocity_action(velocity_size);
    mfem::Vector inactive_displacement_action(velocity_size);
    dv_dv.Mult(velocity_direction, inactive_velocity_action);
    dv_ddisplacement.Mult(displacement_direction, inactive_displacement_action);

    const double density_relative_error  = relative_error(density_finite_difference, density_jacobian_action);
    const double gravity_relative_error  = relative_error(gravity_finite_difference, gravity_jacobian_action);
    const double combined_relative_error = relative_error(combined_finite_difference, combined_jacobian_action);

    INFO("Density Jacobian relative error = " << density_relative_error);
    INFO("Gravity-gradient Jacobian relative error = " << gravity_relative_error);
    INFO("Combined Jacobian relative error = " << combined_relative_error);

    CHECK_THAT(density_relative_error, Catch::Matchers::WithinAbs(0.0, jacobian_tolerance));
    CHECK_THAT(gravity_relative_error, Catch::Matchers::WithinAbs(0.0, jacobian_tolerance));
    CHECK_THAT(combined_relative_error, Catch::Matchers::WithinAbs(0.0, jacobian_tolerance));
    CHECK_THAT(inactive_velocity_action.Norml2(), Catch::Matchers::WithinAbs(0.0, zero_tolerance));
    CHECK_THAT(inactive_displacement_action.Norml2(), Catch::Matchers::WithinAbs(0.0, zero_tolerance));

    mfem::Vector field_coupled_density_action(density_jacobian_action);

    integrator.SetJacobianMode(integrators::GravityForceJacobianMode::minimal);
    dv_dv                                 = 1.0;
    dv_drho                               = 1.0;
    dv_dgrad_phi                          = 1.0;
    dv_ddisplacement                      = 1.0;
    element_state[density_block]          = &density_dofs;
    element_state[gravity_gradient_block] = &gravity_gradient_dofs;
    integrator.AssembleElementGrad(elements, *transformation, element_state, element_matrices);

    mfem::Vector minimal_density_action(velocity_size);
    mfem::Vector minimal_gravity_action(velocity_size);
    dv_drho.Mult(density_direction, minimal_density_action);
    dv_dgrad_phi.Mult(gravity_gradient_direction, minimal_gravity_action);

    const double minimal_density_difference = relative_error(minimal_density_action, field_coupled_density_action);

    INFO("Minimal-mode density-block difference = " << minimal_density_difference);

    CHECK_THAT(minimal_density_difference, Catch::Matchers::WithinAbs(0.0, zero_tolerance));
    CHECK_THAT(minimal_gravity_action.Norml2(), Catch::Matchers::WithinAbs(0.0, zero_tolerance));
}

TEST_CASE(
    "Gravity Force Integrator Matches Manufactured Cartesian Load",
    tags::gravity_integrator_unit
) {
    constexpr int dim                     = 3;
    constexpr double tolerance            = 1.0e-12;

    constexpr int velocity_block          = solver::block_index(solver::FieldBlock::velocity);
    constexpr int density_block           = solver::block_index(solver::FieldBlock::density);
    constexpr int gravity_gradient_block  = solver::block_index(solver::FieldBlock::gravity_gradient);
    constexpr int gravity_potential_block = solver::block_index(solver::FieldBlock::gravity_potential);
    constexpr int displacement_block      = solver::block_index(solver::FieldBlock::displacement);
    constexpr int block_count             = solver::field_block_count;

    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(1, 1, 1, mfem::Element::HEXAHEDRON, 1.0, 1.0, 1.0);

    mfem::H1_FECollection velocity_fec(1, dim);
    mfem::L2_FECollection density_fec(1, dim);
    mfem::RT_FECollection gravity_gradient_fec(0, dim);
    mfem::H1_FECollection displacement_fec(1, dim);
    mfem::H1_FECollection compactification_fec(1, dim);

    mfem::FiniteElementSpace velocity_fes(&mesh, &velocity_fec, dim, mfem::Ordering::byVDIM);
    mfem::FiniteElementSpace density_fes(&mesh, &density_fec);
    mfem::FiniteElementSpace gravity_gradient_fes(&mesh, &gravity_gradient_fec);
    mfem::FiniteElementSpace displacement_fes(&mesh, &displacement_fec, dim, mfem::Ordering::byVDIM);
    mfem::FiniteElementSpace compactification_fes(&mesh, &compactification_fec);

    mfem::GridFunction displacement(&displacement_fes);
    displacement = 0.0;
    mfem::GridFunction compactification_coordinate(&compactification_fes);
    compactification_coordinate         = 0.0;
    mapping::DomainMapper domain_mapper = field_dof_test_utils::make_domain_mapper();

    mapping::GridFunctionMappingEvaluator mapping_evaluator(domain_mapper, displacement, compactification_coordinate);

    auto reference_density          = [](const mfem::Vector &x) { return 1.0 + x(0); };

    auto reference_gravity_gradient = [](const mfem::Vector &x, mfem::Vector &gradient) {
        gradient.SetSize(3);
        gradient(0) = 2.0 * x(0);
        gradient(1) = 3.0 * x(1);
        gradient(2) = 4.0 * x(2);
    };

    mfem::FunctionCoefficient density_coefficient(reference_density);
    mfem::VectorFunctionCoefficient gravity_gradient_coefficient(dim, reference_gravity_gradient);

    mfem::GridFunction density(&density_fes);
    mfem::GridFunction gravity_gradient(&gravity_gradient_fes);
    density.ProjectCoefficient(density_coefficient);
    gravity_gradient.ProjectCoefficient(gravity_gradient_coefficient);

    const mfem::FiniteElement *velocity_element         = velocity_fes.GetFE(0);
    const mfem::FiniteElement *density_element          = density_fes.GetFE(0);
    const mfem::FiniteElement *gravity_gradient_element = gravity_gradient_fes.GetFE(0);
    const mfem::FiniteElement *displacement_element     = displacement_fes.GetFE(0);
    mfem::ElementTransformation *transformation         = mesh.GetElementTransformation(0);

    const int velocity_dofs_count                       = velocity_element->GetDof();
    const int density_dofs_count                        = density_element->GetDof();
    const int gravity_gradient_dofs_count               = gravity_gradient_element->GetDof();
    const int displacement_dofs_count                   = displacement_element->GetDof();
    const int velocity_size                             = dim * velocity_dofs_count;
    const int displacement_size                         = dim * displacement_dofs_count;

    mfem::Array<int> density_dof_indices;
    mfem::Array<int> gravity_gradient_dof_indices;
    mfem::Vector density_dofs;
    mfem::Vector gravity_gradient_dofs;
    density_fes.GetElementDofs(0, density_dof_indices);
    gravity_gradient_fes.GetElementVDofs(0, gravity_gradient_dof_indices);
    density.GetSubVector(density_dof_indices, density_dofs);
    gravity_gradient.GetSubVector(gravity_gradient_dof_indices, gravity_gradient_dofs);

    REQUIRE(density_dofs.Size() == density_dofs_count);
    REQUIRE(gravity_gradient_dofs.Size() == gravity_gradient_dofs_count);

    mfem::Vector velocity_dofs(velocity_size);
    mfem::Vector gravity_potential_dofs(density_dofs_count);
    mfem::Vector displacement_dofs(displacement_size);
    velocity_dofs          = 0.0;
    gravity_potential_dofs = 0.0;
    displacement_dofs      = 0.0;

    mfem::Array<const mfem::FiniteElement *> elements(block_count);
    elements[velocity_block]          = velocity_element;
    elements[density_block]           = density_element;
    elements[gravity_gradient_block]  = gravity_gradient_element;
    elements[gravity_potential_block] = density_element;
    elements[displacement_block]      = displacement_element;

    mfem::Array<const mfem::Vector *> element_state(block_count);
    element_state[velocity_block]          = &velocity_dofs;
    element_state[density_block]           = &density_dofs;
    element_state[gravity_gradient_block]  = &gravity_gradient_dofs;
    element_state[gravity_potential_block] = &gravity_potential_dofs;
    element_state[displacement_block]      = &displacement_dofs;

    mfem::Vector velocity_residual(velocity_size);
    mfem::Vector density_residual(density_dofs_count);
    mfem::Vector gravity_gradient_residual(gravity_gradient_dofs_count);
    mfem::Vector gravity_potential_residual(density_dofs_count);
    mfem::Vector displacement_residual(displacement_size);

    mfem::Array<mfem::Vector *> element_residual(block_count);
    element_residual[velocity_block]          = &velocity_residual;
    element_residual[density_block]           = &density_residual;
    element_residual[gravity_gradient_block]  = &gravity_gradient_residual;
    element_residual[gravity_potential_block] = &gravity_potential_residual;
    element_residual[displacement_block]      = &displacement_residual;

    integrators::GravityMomentumIntegrator integrator(
        domain_mapper, displacement, compactification_coordinate, integrators::GravityForceJacobianMode::field_coupled
    );

    const mfem::IntegrationRule &integration_rule = mfem::IntRules.Get(velocity_element->GetGeomType(), 8);
    integrator.SetIntegrationRule(integration_rule);
    integrator.AssembleElementVector(elements, *transformation, element_state, element_residual);

    mfem::Vector reference_velocity_residual(velocity_residual);

    auto residual_action = [&](const int component, const int coordinate_weight) {
        mfem::Vector test_dofs(velocity_size);
        mfem::Vector x_physical(dim);
        test_dofs                                   = 0.0;

        const mfem::IntegrationRule &velocity_nodes = velocity_element->GetNodes();

        for (int i = 0; i < velocity_dofs_count; ++i) {
            const mfem::IntegrationPoint &node = velocity_nodes.IntPoint(i);
            transformation->SetIntPoint(&node);
            mapping_evaluator.GetPhysicalPoint(*transformation, node, x_physical);
            test_dofs(i + component * velocity_dofs_count) =
                coordinate_weight < 0 ? 1.0 : x_physical(coordinate_weight);
        }

        return test_dofs * velocity_residual;
    };

    CHECK_THAT(residual_action(0, -1), Catch::Matchers::WithinAbs(5.0 / 3.0, tolerance));
    CHECK_THAT(residual_action(1, -1), Catch::Matchers::WithinAbs(9.0 / 4.0, tolerance));
    CHECK_THAT(residual_action(2, -1), Catch::Matchers::WithinAbs(3.0, tolerance));
    CHECK_THAT(residual_action(0, 0), Catch::Matchers::WithinAbs(7.0 / 6.0, tolerance));
    CHECK_THAT(residual_action(1, 1), Catch::Matchers::WithinAbs(3.0 / 2.0, tolerance));
    CHECK_THAT(residual_action(2, 2), Catch::Matchers::WithinAbs(2.0, tolerance));
    CHECK_THAT(residual_action(1, 0), Catch::Matchers::WithinAbs(5.0 / 4.0, tolerance));

    CHECK_THAT(density_residual.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(gravity_gradient_residual.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(gravity_potential_residual.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(displacement_residual.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));

    integrator.SetJacobianMode(integrators::GravityForceJacobianMode::minimal);
    integrator.AssembleElementVector(elements, *transformation, element_state, element_residual);

    mfem::Vector minimal_difference(velocity_residual);
    minimal_difference -= reference_velocity_residual;

    integrator.SetJacobianMode(integrators::GravityForceJacobianMode::exact);
    integrator.AssembleElementVector(elements, *transformation, element_state, element_residual);

    mfem::Vector exact_difference(velocity_residual);
    exact_difference -= reference_velocity_residual;

    CHECK_THAT(minimal_difference.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(exact_difference.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
}
TEST_CASE(
    "Gravity Force Integrator Preserves Gravity Identities",
    tags::gravity_integrator_unit
) {
    constexpr int dim                     = 3;
    constexpr double density_value        = 1.7;
    constexpr double gravity_scale        = 2.4;
    constexpr double density_scale        = 0.6;
    constexpr double tolerance            = 1.0e-12;

    constexpr int velocity_block          = solver::block_index(solver::FieldBlock::velocity);
    constexpr int density_block           = solver::block_index(solver::FieldBlock::density);
    constexpr int gravity_gradient_block  = solver::block_index(solver::FieldBlock::gravity_gradient);
    constexpr int gravity_potential_block = solver::block_index(solver::FieldBlock::gravity_potential);
    constexpr int displacement_block      = solver::block_index(solver::FieldBlock::displacement);
    constexpr int block_count             = solver::field_block_count;

    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(1, 1, 1, mfem::Element::HEXAHEDRON, 1.0, 1.0, 1.0);

    mfem::H1_FECollection velocity_fec(1, dim);
    mfem::L2_FECollection density_fec(0, dim);
    mfem::RT_FECollection gravity_gradient_fec(0, dim);
    mfem::H1_FECollection displacement_fec(1, dim);
    mfem::H1_FECollection compactification_fec(1, dim);

    mfem::FiniteElementSpace velocity_fes(&mesh, &velocity_fec, dim, mfem::Ordering::byVDIM);
    mfem::FiniteElementSpace density_fes(&mesh, &density_fec);
    mfem::FiniteElementSpace gravity_gradient_fes(&mesh, &gravity_gradient_fec);
    mfem::FiniteElementSpace displacement_fes(&mesh, &displacement_fec, dim, mfem::Ordering::byVDIM);
    mfem::FiniteElementSpace compactification_fes(&mesh, &compactification_fec);

    mfem::GridFunction displacement(&displacement_fes);
    displacement = 0.0;
    mfem::GridFunction compactification_coordinate(&compactification_fes);
    compactification_coordinate         = 0.0;
    mapping::DomainMapper domain_mapper = field_dof_test_utils::make_domain_mapper();
    mapping::GridFunctionMappingEvaluator mapping_evaluator(domain_mapper, displacement, compactification_coordinate);

    auto radial_gravity = [](const mfem::Vector &x, mfem::Vector &gravity) {
        gravity.SetSize(3);
        gravity(0) = x(0) - 0.5;
        gravity(1) = x(1) - 0.5;
        gravity(2) = x(2) - 0.5;
    };

    mfem::ConstantCoefficient density_coefficient(density_value);
    mfem::VectorFunctionCoefficient gravity_coefficient(dim, radial_gravity);

    mfem::GridFunction density(&density_fes);
    mfem::GridFunction gravity_gradient(&gravity_gradient_fes);
    density.ProjectCoefficient(density_coefficient);
    gravity_gradient.ProjectCoefficient(gravity_coefficient);

    const mfem::FiniteElement *velocity_element         = velocity_fes.GetFE(0);
    const mfem::FiniteElement *density_element          = density_fes.GetFE(0);
    const mfem::FiniteElement *gravity_gradient_element = gravity_gradient_fes.GetFE(0);
    const mfem::FiniteElement *displacement_element     = displacement_fes.GetFE(0);
    mfem::ElementTransformation *transformation         = mesh.GetElementTransformation(0);

    const int velocity_dofs_count                       = velocity_element->GetDof();
    const int density_dofs_count                        = density_element->GetDof();
    const int gravity_gradient_dofs_count               = gravity_gradient_element->GetDof();
    const int displacement_dofs_count                   = displacement_element->GetDof();
    const int velocity_size                             = dim * velocity_dofs_count;
    const int displacement_size                         = dim * displacement_dofs_count;

    mfem::Array<int> density_dof_indices;
    mfem::Array<int> gravity_gradient_dof_indices;
    mfem::Vector density_dofs;
    mfem::Vector gravity_gradient_dofs;
    density_fes.GetElementDofs(0, density_dof_indices);
    gravity_gradient_fes.GetElementVDofs(0, gravity_gradient_dof_indices);
    density.GetSubVector(density_dof_indices, density_dofs);
    gravity_gradient.GetSubVector(gravity_gradient_dof_indices, gravity_gradient_dofs);

    mfem::Vector zero_density(density_dofs_count);
    mfem::Vector zero_gravity(gravity_gradient_dofs_count);
    mfem::Vector velocity_dofs(velocity_size);
    mfem::Vector gravity_potential_dofs(density_dofs_count);
    mfem::Vector displacement_dofs(displacement_size);
    zero_density           = 0.0;
    zero_gravity           = 0.0;
    velocity_dofs          = 0.0;
    gravity_potential_dofs = 0.0;
    displacement_dofs      = 0.0;

    mfem::Array<const mfem::FiniteElement *> elements(block_count);
    elements[velocity_block]          = velocity_element;
    elements[density_block]           = density_element;
    elements[gravity_gradient_block]  = gravity_gradient_element;
    elements[gravity_potential_block] = density_element;
    elements[displacement_block]      = displacement_element;

    mfem::Array<const mfem::Vector *> element_state(block_count);
    element_state[velocity_block]          = &velocity_dofs;
    element_state[density_block]           = &density_dofs;
    element_state[gravity_gradient_block]  = &gravity_gradient_dofs;
    element_state[gravity_potential_block] = &gravity_potential_dofs;
    element_state[displacement_block]      = &displacement_dofs;

    mfem::Vector velocity_residual(velocity_size);
    mfem::Vector density_residual(density_dofs_count);
    mfem::Vector gravity_gradient_residual(gravity_gradient_dofs_count);
    mfem::Vector gravity_potential_residual(density_dofs_count);
    mfem::Vector displacement_residual(displacement_size);

    mfem::Array<mfem::Vector *> element_residual(block_count);
    element_residual[velocity_block]          = &velocity_residual;
    element_residual[density_block]           = &density_residual;
    element_residual[gravity_gradient_block]  = &gravity_gradient_residual;
    element_residual[gravity_potential_block] = &gravity_potential_residual;
    element_residual[displacement_block]      = &displacement_residual;

    integrators::GravityMomentumIntegrator integrator(
        domain_mapper, displacement, compactification_coordinate, integrators::GravityForceJacobianMode::field_coupled
    );

    const mfem::IntegrationRule &integration_rule = mfem::IntRules.Get(velocity_element->GetGeomType(), 8);
    integrator.SetIntegrationRule(integration_rule);

    auto assemble_velocity_residual = [&](const mfem::Vector &density_state, const mfem::Vector &gravity_state) {
        element_state[density_block]          = &density_state;
        element_state[gravity_gradient_block] = &gravity_state;
        integrator.AssembleElementVector(elements, *transformation, element_state, element_residual);
        return mfem::Vector(velocity_residual);
    };

    auto scaled_difference_norm = [](mfem::Vector computed, const mfem::Vector &reference, const double scale) {
        computed.Add(-scale, reference);
        return computed.Norml2();
    };

    const mfem::Vector base_residual = assemble_velocity_residual(density_dofs, gravity_gradient_dofs);

    REQUIRE(base_residual.Norml2() > tolerance);

    const mfem::Vector zero_field_residual = assemble_velocity_residual(density_dofs, zero_gravity);

    mfem::Vector reversed_gravity(gravity_gradient_dofs);
    reversed_gravity *= -1.0;
    const mfem::Vector reversed_residual = assemble_velocity_residual(density_dofs, reversed_gravity);

    mfem::Vector scaled_gravity(gravity_gradient_dofs);
    scaled_gravity *= gravity_scale;
    const mfem::Vector gravity_scaled_residual = assemble_velocity_residual(density_dofs, scaled_gravity);

    mfem::Vector scaled_density(density_dofs);
    scaled_density *= density_scale;
    const mfem::Vector density_scaled_residual = assemble_velocity_residual(scaled_density, gravity_gradient_dofs);
    const mfem::Vector jointly_scaled_residual = assemble_velocity_residual(scaled_density, scaled_gravity);

    CHECK_THAT(zero_field_residual.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(
        scaled_difference_norm(reversed_residual, base_residual, -1.0), Catch::Matchers::WithinAbs(0.0, tolerance)
    );
    CHECK_THAT(
        scaled_difference_norm(gravity_scaled_residual, base_residual, gravity_scale),
        Catch::Matchers::WithinAbs(0.0, tolerance)
    );
    CHECK_THAT(
        scaled_difference_norm(density_scaled_residual, base_residual, density_scale),
        Catch::Matchers::WithinAbs(0.0, tolerance)
    );
    CHECK_THAT(
        scaled_difference_norm(jointly_scaled_residual, base_residual, density_scale * gravity_scale),
        Catch::Matchers::WithinAbs(0.0, tolerance)
    );

    auto make_test_dofs = [&](auto &&test_function) {
        mfem::Vector test_dofs(velocity_size);
        mfem::Vector x_physical(dim);
        mfem::Vector centered_position(dim);
        mfem::Vector test_value(dim);
        test_dofs                                   = 0.0;

        const mfem::IntegrationRule &velocity_nodes = velocity_element->GetNodes();

        for (int i = 0; i < velocity_dofs_count; ++i) {
            const mfem::IntegrationPoint &node = velocity_nodes.IntPoint(i);
            transformation->SetIntPoint(&node);
            mapping_evaluator.GetPhysicalPoint(*transformation, node, x_physical);

            for (int component = 0; component < dim; ++component) {
                centered_position(component) = x_physical(component) - 0.5;
            }

            test_function(centered_position, test_value);

            for (int component = 0; component < dim; ++component) {
                test_dofs(i + component * velocity_dofs_count) = test_value(component);
            }
        }

        return test_dofs;
    };

    const mfem::Vector force_x_test  = make_test_dofs([](const mfem::Vector &, mfem::Vector &value) {
        value.SetSize(3);
        value    = 0.0;
        value(0) = 1.0;
    });

    const mfem::Vector force_y_test  = make_test_dofs([](const mfem::Vector &, mfem::Vector &value) {
        value.SetSize(3);
        value    = 0.0;
        value(1) = 1.0;
    });

    const mfem::Vector force_z_test  = make_test_dofs([](const mfem::Vector &, mfem::Vector &value) {
        value.SetSize(3);
        value    = 0.0;
        value(2) = 1.0;
    });

    const mfem::Vector torque_x_test = make_test_dofs([](const mfem::Vector &position, mfem::Vector &value) {
        value.SetSize(3);
        value(0) = 0.0;
        value(1) = -position(2);
        value(2) = position(1);
    });

    const mfem::Vector torque_y_test = make_test_dofs([](const mfem::Vector &position, mfem::Vector &value) {
        value.SetSize(3);
        value(0) = position(2);
        value(1) = 0.0;
        value(2) = -position(0);
    });

    const mfem::Vector torque_z_test = make_test_dofs([](const mfem::Vector &position, mfem::Vector &value) {
        value.SetSize(3);
        value(0) = -position(1);
        value(1) = position(0);
        value(2) = 0.0;
    });

    CHECK_THAT(force_x_test * base_residual, Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(force_y_test * base_residual, Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(force_z_test * base_residual, Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(torque_x_test * base_residual, Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(torque_y_test * base_residual, Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(torque_z_test * base_residual, Catch::Matchers::WithinAbs(0.0, tolerance));

    mfem::DenseMatrix dv_drho(velocity_size, density_dofs_count);
    mfem::DenseMatrix dv_dgrad_phi(velocity_size, gravity_gradient_dofs_count);

    mfem::Array2D<mfem::DenseMatrix *> element_matrices(block_count, block_count);

    for (int row = 0; row < block_count; ++row) {
        for (int column = 0; column < block_count; ++column) {
            element_matrices(row, column) = nullptr;
        }
    }

    element_matrices(velocity_block, density_block)          = &dv_drho;
    element_matrices(velocity_block, gravity_gradient_block) = &dv_dgrad_phi;

    element_state[density_block]                             = &density_dofs;
    element_state[gravity_gradient_block]                    = &zero_gravity;
    integrator.AssembleElementGrad(elements, *transformation, element_state, element_matrices);

    mfem::Vector zero_gravity_density_action(velocity_size);
    mfem::Vector zero_gravity_field_action(velocity_size);
    dv_drho.Mult(density_dofs, zero_gravity_density_action);
    dv_dgrad_phi.Mult(gravity_gradient_dofs, zero_gravity_field_action);

    CHECK_THAT(zero_gravity_density_action.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
    CHECK_THAT(
        scaled_difference_norm(zero_gravity_field_action, base_residual, 1.0),
        Catch::Matchers::WithinAbs(0.0, tolerance)
    );

    element_state[density_block]          = &zero_density;
    element_state[gravity_gradient_block] = &gravity_gradient_dofs;
    integrator.AssembleElementGrad(elements, *transformation, element_state, element_matrices);

    mfem::Vector zero_density_density_action(velocity_size);
    mfem::Vector zero_density_field_action(velocity_size);
    dv_drho.Mult(density_dofs, zero_density_density_action);
    dv_dgrad_phi.Mult(gravity_gradient_dofs, zero_density_field_action);

    CHECK_THAT(
        scaled_difference_norm(zero_density_density_action, base_residual, 1.0),
        Catch::Matchers::WithinAbs(0.0, tolerance)
    );
    CHECK_THAT(zero_density_field_action.Norml2(), Catch::Matchers::WithinAbs(0.0, tolerance));
}