#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::unit &tags::solver &tags::integrator &tags::gravity
) {
    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::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::GridFunction displacement(&displacement_fes);
    displacement = 0.0;

    mapping::DomainMapper domain_mapper(displacement, 1.0, 2.0);
    INFO(
        std::format(
            "Domain mapping is has displacement field: {}",
            domain_mapper.HasDisplacementField()
        )
    );
    INFO(
        std::format(
            "Domain mapping is identity: {}", domain_mapper.CalcIsIdentity()
        )
    );

    REQUIRE(domain_mapper.CalcIsIdentity());

    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, 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::unit &tags::solver &tags::integrator &tags::gravity
) {
    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::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::GridFunction displacement(&displacement_fes);
    displacement = 0.0;

    mapping::DomainMapper domain_mapper(displacement, 1.0, 2.0);

    REQUIRE(domain_mapper.CalcIsIdentity());

    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, 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);
            domain_mapper.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::unit &tags::solver &tags::integrator &tags::gravity
) {
    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::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::GridFunction displacement(&displacement_fes);
    displacement = 0.0;

    mapping::DomainMapper domain_mapper(displacement, 1.0, 2.0);

    REQUIRE(domain_mapper.HasDisplacementField());

    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, 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);
            domain_mapper.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)
    );
}