#include <algorithm>
#include <array>
#include <cmath>
#include <limits>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace gravity_displacement_force_test_utils {
    using CoupledForm = mean_field::utils::blocks::barotropic_equilibrium_form;

    constexpr auto densityValue =
        mean_field::utils::blocks::get_value_block<CoupledForm>(mean_field::utils::blocks::density_field.mass_term);

    constexpr auto displacementValue = mean_field::utils::blocks::get_value_block<CoupledForm>(
        mean_field::utils::blocks::displacement_field.geometry_term
    );

    constexpr auto gravityGradientValue =
        mean_field::utils::blocks::get_value_block<CoupledForm>(mean_field::utils::blocks::gravity_field.gradient_term);

    constexpr auto gravityPotentialValue =
        mean_field::utils::blocks::get_value_block<CoupledForm>(mean_field::utils::blocks::gravity_field.poisson_term);

    constexpr auto enthalpyValue = mean_field::utils::blocks::get_value_block<CoupledForm>(
        mean_field::utils::blocks::enthalpy_field.specific_term
    );

    constexpr auto barotropicConstantValue = mean_field::utils::blocks::get_value_block<CoupledForm>(
        mean_field::utils::blocks::barotropic_constant_field.mass_normalization_term
    );

    constexpr auto gravityGradientResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::gravity_field.gradient_term
    );

    constexpr auto gravityPotentialResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::gravity_field.poisson_term
    );

    constexpr auto densityResidual =
        mean_field::utils::blocks::get_residual_block<CoupledForm>(mean_field::utils::blocks::density_field.mass_term);

    constexpr auto displacementResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::displacement_field.geometry_term
    );

    constexpr auto enthalpyResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::enthalpy_field.specific_term
    );

    constexpr auto massResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::barotropic_constant_field.mass_normalization_term
    );

    [[nodiscard]] mean_field::operators::GravityDisplacementForceLayout make_layout(const mean_field::fem::FEM &f) {
        using DomainSchema         = gravity_prepared_test_utils::DomainSchema;

        const auto densityMap      = gravity_prepared_test_utils::make_field_map<mean_field::field::Density>(f);
        const auto displacementMap = gravity_prepared_test_utils::make_field_map<mean_field::field::Displacement>(f);
        const auto gravityFluxMap =
            mean_field::field::make_field_dof_map<mean_field::field::Gravity, DomainSchema>(*f.gravityFluxFes);
        const auto gravityPotentialMap =
            mean_field::field::make_field_dof_map<mean_field::field::Gravity, DomainSchema>(*f.gravityPotentialFes);
        const auto enthalpyMap =
            mean_field::field::make_field_dof_map<mean_field::field::Enthalpy, DomainSchema>(*f.enthalpyFes);

        const std::array<int, CoupledForm::value_block_count> valueSizes{
            densityMap.reduced_size(),          displacementMap.reduced_size(), gravityFluxMap.reduced_size(),
            gravityPotentialMap.reduced_size(), enthalpyMap.reduced_size(),     1
        };

        const std::array<int, CoupledForm::residual_block_count> residualSizes{
            gravityFluxMap.reduced_size(),  gravityPotentialMap.reduced_size(), densityMap.reduced_size(),
            displacementMap.reduced_size(), enthalpyMap.reduced_size(),         1
        };

        return {valueSizes, residualSizes};
    }

    [[nodiscard]] mfem::Vector make_density(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction densityField(f.densityFes.get());

        mfem::FunctionCoefficient densityCoefficient([phase](const mfem::Vector &position) {
            return 0.82 + 0.07 * std::sin(0.8 * position(0) + phase) +
                   0.05 * std::cos(0.6 * position(1) - 0.3 * phase) + 0.03 * position(2) * position(2);
        });

        densityField.ProjectCoefficient(densityCoefficient);

        mfem::Vector densityTrue;
        densityField.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] mfem::Vector make_density_direction(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction densityField(f.densityFes.get());

        mfem::FunctionCoefficient densityCoefficient([phase](const mfem::Vector &position) {
            return 0.19 * std::sin(0.9 * position(0) + phase) - 0.13 * std::cos(0.7 * position(1) - phase) +
                   0.08 * position(2);
        });

        densityField.ProjectCoefficient(densityCoefficient);

        mfem::Vector densityTrue;
        densityField.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] mfem::Vector make_gravity_gradient(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction gravityField(f.gravityFluxFes.get());

        auto gravityFunction = [phase](const mfem::Vector &position, mfem::Vector &value) {
            value.SetSize(3);

            value(0) = 0.31 + 0.08 * position(0) + 0.03 * phase * position(1);

            value(1) = -0.17 + 0.06 * position(1) - 0.02 * phase * position(2);

            value(2) = 0.23 - 0.05 * position(2) + 0.025 * phase * position(0);
        };

        mfem::VectorFunctionCoefficient gravityCoefficient(3, gravityFunction);

        gravityField.ProjectCoefficient(gravityCoefficient);

        mfem::Vector gravityTrue;
        gravityField.GetTrueDofs(gravityTrue);
        return gravityTrue;
    }

    [[nodiscard]] mfem::Vector make_gravity_gradient_direction(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction gravityField(f.gravityFluxFes.get());

        auto gravityFunction = [phase](const mfem::Vector &position, mfem::Vector &value) {
            value.SetSize(3);

            value(0) = 0.14 * std::sin(position(0) + phase) + 0.03 * position(1);

            value(1) = -0.11 * std::cos(position(1) - phase) + 0.04 * position(2);

            value(2) = 0.09 * std::sin(position(2) + 0.5 * phase) - 0.02 * position(0);
        };

        mfem::VectorFunctionCoefficient gravityCoefficient(3, gravityFunction);

        gravityField.ProjectCoefficient(gravityCoefficient);

        mfem::Vector gravityTrue;
        gravityField.GetTrueDofs(gravityTrue);
        return gravityTrue;
    }

    [[nodiscard]] mfem::Vector make_displacement_direction(const mean_field::fem::FEM &f) {
        mfem::Vector direction    = gravity_prepared_test_utils::make_displacement(f, 0.83);

        const mfem::Vector second = gravity_prepared_test_utils::make_displacement(f, 0.29);

        direction -= second;
        return direction;
    }

    [[nodiscard]] mfem::Vector make_vacuum_only_density(const mean_field::fem::FEM &f) {
        mfem::ParGridFunction densityField(f.densityFes.get());
        densityField              = 0.0;

        const int vacuumAttribute = field_dof_test_utils::vacuum_material_attribute;

        mfem::Array<int> densityDofs;
        int localVacuumElements = 0;

        for (int elementId = 0; elementId < f.mesh->GetNE(); ++elementId) {
            mfem::ElementTransformation *transformation = f.mesh->GetElementTransformation(elementId);

            REQUIRE(transformation != nullptr);

            if (transformation->Attribute != vacuumAttribute) {
                continue;
            }

            f.densityFes->GetElementDofs(elementId, densityDofs);

            mfem::Vector elementDensity(densityDofs.Size());
            elementDensity = 1.0;
            densityField.SetSubVector(densityDofs, elementDensity);
            ++localVacuumElements;
        }

        int globalVacuumElements = 0;
        MPI_Allreduce(&localVacuumElements, &globalVacuumElements, 1, MPI_INT, MPI_SUM, f.mesh->GetComm());

        REQUIRE(globalVacuumElements > 0);

        mfem::Vector densityTrue;
        densityField.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] mean_field::operators::context::gravity_field::GravityFieldRevisions make_revisions() {
        return {
            .discretization    = {.value = 3},
            .displacement      = {.value = 5},
            .density           = {.value = 7},
            .gravity_gradient  = {.value = 11},
            .gravity_potential = {.value = 13}
        };
    }

    void prepare_gravity_context(
        mean_field::operators::context::gravity_field::GravityFieldLinearizationContext &context,
        const mfem::Vector &density,
        const mfem::Vector &displacement,
        const mfem::Vector &gravityGradient,
        const mfem::Vector &gravityPotential,
        const mean_field::operators::context::gravity_field::GravityFieldRevisions &revisions
    ) {
        context.Prepare(
            {.density           = context.GetDensityMap().gather(density),
             .displacement      = context.GetDisplacementMap().gather(displacement),
             .gravity_gradient  = context.GetGravityGradientMap().gather(gravityGradient),
             .gravity_potential = context.GetGravityPotentialMap().gather(gravityPotential)},
            revisions
        );
    }

    [[nodiscard]] double relative_difference(
        const mfem::Vector &left,
        const mfem::Vector &right,
        const MPI_Comm communicator
    ) {
        MFEM_VERIFY(
            left.Size() == right.Size(), "Cannot compare gravity-displacement-force vectors with "
                                         "different sizes."
        );

        mfem::Vector difference(left);
        difference -= right;

        const double scale = std::max(
            {gravity_prepared_test_utils::global_norm(left, communicator),
             gravity_prepared_test_utils::global_norm(right, communicator),
             100.0 * std::numeric_limits<double>::epsilon()}
        );

        return gravity_prepared_test_utils::global_norm(difference, communicator) / scale;
    }

    [[nodiscard]] mfem::Vector centered_difference(
        const mean_field::fem::FEM &f,
        const mfem::Vector &baseDensity,
        const mfem::Vector &densityDirection,
        const mfem::Vector &baseGravityGradient,
        const mfem::Vector &gravityGradientDirection,
        const mfem::Vector &baseDisplacement,
        const mfem::Vector &displacementDirection,
        const double step
    ) {
        mfem::Vector plusDensity(baseDensity);
        plusDensity.Add(step, densityDirection);

        mfem::Vector minusDensity(baseDensity);
        minusDensity.Add(-step, densityDirection);

        mfem::Vector plusGravity(baseGravityGradient);
        plusGravity.Add(step, gravityGradientDirection);

        mfem::Vector minusGravity(baseGravityGradient);
        minusGravity.Add(-step, gravityGradientDirection);

        mfem::Vector plusDisplacement(baseDisplacement);
        plusDisplacement.Add(step, displacementDirection);

        mfem::Vector minusDisplacement(baseDisplacement);
        minusDisplacement.Add(-step, displacementDirection);

        mfem::Vector plusResidual;
        mfem::Vector minusResidual;

        mean_field::operators::kernels::apply_gravity_displacement_force_residual(
            f, *f.domainMapperStateless, plusDensity, plusGravity, plusDisplacement, plusResidual
        );

        mean_field::operators::kernels::apply_gravity_displacement_force_residual(
            f, *f.domainMapperStateless, minusDensity, minusGravity, minusDisplacement, minusResidual
        );

        plusResidual -= minusResidual;
        plusResidual /= 2.0 * step;
        return plusResidual;
    }

    template <int index>
    [[nodiscard]] mfem::Vector copy_residual_block(
        const mfem::Vector &action,
        const mean_field::operators::GravityDisplacementForceLayout &layout,
        const mean_field::utils::blocks::residual_block<index> block
    ) {
        mfem::Vector result(layout.size(block));
        const int offset = layout.offset(block);

        for (int entry = 0; entry < result.Size(); ++entry) {
            result(entry) = action(offset + entry);
        }

        return result;
    }
} // namespace gravity_displacement_force_test_utils

TEST_CASE(
    "Gravity Displacement Force Query Includes Every Registered Operand",
    tags::gravity_unit
) {
    using DisplacementField           = mean_field::field::Field<mean_field::field::Displacement>;

    constexpr int geometryWeightOrder = 4;

    constexpr mean_field::quadrature::Query query =
        DisplacementField::make_query<mean_field::field::Displacement::Form::GravityForce>(
            mean_field::quadrature::QuadratureRole::discretization, geometryWeightOrder, {},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

    /*
     * rho: 2
     * RT value: family order 2 + 1 = 3
     * geometry displacement gradient: 3 - 1 = 2
     * displacement test value: 3
     * reference-element geometry weight: 4
     */
    constexpr int expectedBaseOrder = 2 + 3 + 2 + 3 + 4;

    STATIC_REQUIRE(query.term == mean_field::quadrature::Term::gravity_force);

    STATIC_REQUIRE(query.role == mean_field::quadrature::QuadratureRole::discretization);

    STATIC_REQUIRE(query.domain == mean_field::utils::DOMAINS::STELLAR);

    STATIC_REQUIRE(query.mapping == mean_field::quadrature::MappingKind::general);

    STATIC_REQUIRE(query.base_order.has_value());
    STATIC_REQUIRE(*query.base_order == expectedBaseOrder);
}

TEST_CASE(
    "Gravity Displacement Force Uses Positive Grad-Phi Sign And Excludes "
    "Vacuum",
    tags::gravity_kernel_accuracy
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    mfem::ParGridFunction densityField(f.densityFes.get());
    mfem::ConstantCoefficient densityCoefficient(1.0);
    densityField.ProjectCoefficient(densityCoefficient);

    mfem::Vector density;
    densityField.GetTrueDofs(density);

    mfem::ParGridFunction gravityField(f.gravityFluxFes.get());

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

    mfem::VectorFunctionCoefficient gravityCoefficient(3, constantGravityFunction);

    gravityField.ProjectCoefficient(gravityCoefficient);

    mfem::Vector gravityGradient;
    gravityField.GetTrueDofs(gravityGradient);

    mfem::Vector displacement(f.displacementFes->GetTrueVSize());
    displacement = 0.0;

    mfem::Vector residual;

    mean_field::operators::kernels::apply_gravity_displacement_force_residual(
        f, *f.domainMapperStateless, density, gravityGradient, displacement, residual
    );

    mfem::ParGridFunction testField(f.displacementFes.get());
    testField.ProjectCoefficient(gravityCoefficient);

    mfem::Vector testDirection;
    testField.GetTrueDofs(testDirection);

    const double signedWork = gravity_prepared_test_utils::global_dot(residual, testDirection, f.mesh->GetComm());

    INFO("Constant +x gravity-force work = " << signedWork);
    CHECK(signedWork > 0.0);

    const mfem::Vector vacuumDensity = gravity_displacement_force_test_utils::make_vacuum_only_density(f);

    mfem::Vector vacuumResidual;

    mean_field::operators::kernels::apply_gravity_displacement_force_residual(
        f, *f.domainMapperStateless, vacuumDensity, gravityGradient, displacement, vacuumResidual
    );

    CHECK(gravity_prepared_test_utils::global_norm(vacuumResidual, f.mesh->GetComm()) == 0.0);
}

TEST_CASE(
    "Prepared Gravity Displacement Force Reuses Shared Gravity Revisions",
    tags::gravity_prepared
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    mfem::Vector density               = gravity_displacement_force_test_utils::make_density(f, 0.31);

    const mfem::Vector gravityGradient = gravity_displacement_force_test_utils::make_gravity_gradient(f, 0.47);

    const mfem::Vector displacement    = gravity_prepared_test_utils::make_displacement(f, 0.61);

    mfem::Vector gravityPotential(f.gravityPotentialFes->GetTrueVSize());
    gravityPotential = 0.0;

    auto revisions   = gravity_displacement_force_test_utils::make_revisions();

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );

    gravity_displacement_force_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, revisions
    );

    mean_field::operators::PreparedGravityDisplacementForceOperator preparedOperator(
        f, *f.domainMapperStateless, gravityContext
    );

    const auto initialReport = preparedOperator.Prepare();
    REQUIRE(initialReport.DidAnyWork());
    REQUIRE(preparedOperator.IsPrepared());

    mfem::Vector preparedResidual;
    mfem::Vector kernelResidual;

    preparedOperator.BuildResidual(preparedResidual);

    mean_field::operators::kernels::apply_gravity_displacement_force_residual(
        f, *f.domainMapperStateless, density, gravityGradient, displacement, kernelResidual
    );

    const mfem::Vector kernelResidualReduced = gravityContext.GetDisplacementMap().gather(kernelResidual);

    CHECK(
        gravity_displacement_force_test_utils::relative_difference(
            preparedResidual, kernelResidualReduced, f.mesh->GetComm()
        ) < 2.0e-12
    );

    CHECK_FALSE(preparedOperator.Prepare().DidAnyWork());

    ++revisions.gravity_potential.value;

    gravity_displacement_force_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, revisions
    );

    CHECK(preparedOperator.IsPrepared());
    CHECK_FALSE(preparedOperator.Prepare().DidAnyWork());

    density = gravity_displacement_force_test_utils::make_density(f, 0.79);
    ++revisions.density.value;

    gravity_displacement_force_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, revisions
    );

    CHECK_FALSE(preparedOperator.IsPrepared());

    const auto densityReport = preparedOperator.Prepare();
    CHECK(densityReport.DidAnyWork());
    CHECK(preparedOperator.IsPrepared());
    CHECK(preparedOperator.GetResidualPreparationCount() == 2);
    CHECK(preparedOperator.GetResidualApplicationCount() == 1);
}

TEST_CASE(
    "Gravity Displacement Force Jacobian Matches All Columns And Centered "
    "Differences",
    tags::gravity_prepared_jacobian_accuracy
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    const mfem::Vector density          = gravity_displacement_force_test_utils::make_density(f, 0.37);

    const mfem::Vector densityDirection = gravity_displacement_force_test_utils::make_density_direction(f, 0.53);

    const mfem::Vector gravityGradient  = gravity_displacement_force_test_utils::make_gravity_gradient(f, 0.67);

    const mfem::Vector gravityGradientDirection =
        gravity_displacement_force_test_utils::make_gravity_gradient_direction(f, 0.71);

    const mfem::Vector displacement          = gravity_prepared_test_utils::make_displacement(f, 0.59);

    const mfem::Vector displacementDirection = gravity_displacement_force_test_utils::make_displacement_direction(f);

    mfem::Vector gravityPotential(f.gravityPotentialFes->GetTrueVSize());
    gravityPotential = 0.0;

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );

    gravity_displacement_force_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential,
        gravity_displacement_force_test_utils::make_revisions()
    );

    mean_field::operators::PreparedGravityDisplacementForceOperator preparedOperator(
        f, *f.domainMapperStateless, gravityContext
    );

    preparedOperator.Prepare();

    const mfem::Vector densityDirectionReduced = gravityContext.GetDensityMap().gather(densityDirection);
    const mfem::Vector gravityGradientDirectionReduced =
        gravityContext.GetGravityGradientMap().gather(gravityGradientDirection);
    const mfem::Vector displacementDirectionReduced = gravityContext.GetDisplacementMap().gather(displacementDirection);

    mfem::Vector densityAction;
    mfem::Vector gravityAction;
    mfem::Vector displacementAction;
    mfem::Vector completeAction;

    preparedOperator.ApplyDensityJacobianAction(densityDirectionReduced, densityAction);

    preparedOperator.ApplyGravityGradientJacobianAction(gravityGradientDirectionReduced, gravityAction);

    preparedOperator.ApplyDisplacementJacobianAction(displacementDirectionReduced, displacementAction);

    preparedOperator.ApplyCompleteJacobianAction(
        densityDirectionReduced, displacementDirectionReduced, gravityGradientDirectionReduced, completeAction
    );

    mfem::Vector summedColumns(densityAction);
    summedColumns += gravityAction;
    summedColumns += displacementAction;

    CHECK(
        gravity_displacement_force_test_utils::relative_difference(completeAction, summedColumns, f.mesh->GetComm()) <
        2.0e-12
    );

    mfem::Vector zeroDensity(densityDirection.Size());
    mfem::Vector zeroGravity(gravityGradientDirection.Size());
    mfem::Vector zeroDisplacement(displacementDirection.Size());
    zeroDensity                              = 0.0;
    zeroGravity                              = 0.0;
    zeroDisplacement                         = 0.0;

    constexpr double step                    = 1.0e-5;

    const mfem::Vector densityDifferenceTrue = gravity_displacement_force_test_utils::centered_difference(
        f, density, densityDirection, gravityGradient, zeroGravity, displacement, zeroDisplacement, step
    );

    const mfem::Vector gravityDifferenceTrue = gravity_displacement_force_test_utils::centered_difference(
        f, density, zeroDensity, gravityGradient, gravityGradientDirection, displacement, zeroDisplacement, step
    );

    const mfem::Vector displacementDifferenceTrue = gravity_displacement_force_test_utils::centered_difference(
        f, density, zeroDensity, gravityGradient, zeroGravity, displacement, displacementDirection, step
    );

    const mfem::Vector completeDifferenceTrue = gravity_displacement_force_test_utils::centered_difference(
        f, density, densityDirection, gravityGradient, gravityGradientDirection, displacement, displacementDirection,
        step
    );

    const mfem::Vector densityDifference      = gravityContext.GetDisplacementMap().gather(densityDifferenceTrue);
    const mfem::Vector gravityDifference      = gravityContext.GetDisplacementMap().gather(gravityDifferenceTrue);
    const mfem::Vector displacementDifference = gravityContext.GetDisplacementMap().gather(displacementDifferenceTrue);
    const mfem::Vector completeDifference     = gravityContext.GetDisplacementMap().gather(completeDifferenceTrue);

    const double densityError =
        gravity_displacement_force_test_utils::relative_difference(densityAction, densityDifference, f.mesh->GetComm());

    const double gravityError =
        gravity_displacement_force_test_utils::relative_difference(gravityAction, gravityDifference, f.mesh->GetComm());

    const double displacementError = gravity_displacement_force_test_utils::relative_difference(
        displacementAction, displacementDifference, f.mesh->GetComm()
    );

    const double completeError = gravity_displacement_force_test_utils::relative_difference(
        completeAction, completeDifference, f.mesh->GetComm()
    );

    INFO("Density-column centered-difference error = " << densityError);
    INFO("Gravity-column centered-difference error = " << gravityError);
    INFO("Displacement-column centered-difference error = " << displacementError);
    INFO("Complete centered-difference error = " << completeError);

    CHECK(densityError < 2.0e-9);
    CHECK(gravityError < 2.0e-9);
    CHECK(displacementError < 2.0e-8);
    CHECK(completeError < 3.0e-8);
}

TEST_CASE(
    "Prepared Gravity Displacement Force MFEM Adapter Routes Only R-d",
    tags::gravity_prepared_unit
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    const mfem::Vector density          = gravity_displacement_force_test_utils::make_density(f, 0.41);

    const mfem::Vector densityDirection = gravity_displacement_force_test_utils::make_density_direction(f, 0.57);

    const mfem::Vector gravityGradient  = gravity_displacement_force_test_utils::make_gravity_gradient(f, 0.63);

    const mfem::Vector gravityGradientDirection =
        gravity_displacement_force_test_utils::make_gravity_gradient_direction(f, 0.77);

    const mfem::Vector displacement          = gravity_prepared_test_utils::make_displacement(f, 0.51);

    const mfem::Vector displacementDirection = gravity_displacement_force_test_utils::make_displacement_direction(f);

    mfem::Vector gravityPotential(f.gravityPotentialFes->GetTrueVSize());
    gravityPotential = 0.0;

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );

    gravity_displacement_force_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential,
        gravity_displacement_force_test_utils::make_revisions()
    );

    mean_field::operators::PreparedGravityDisplacementForceOperator preparedOperator(
        f, *f.domainMapperStateless, gravityContext
    );

    preparedOperator.Prepare();

    const mfem::Vector densityDirectionReduced = gravityContext.GetDensityMap().gather(densityDirection);
    const mfem::Vector gravityGradientDirectionReduced =
        gravityContext.GetGravityGradientMap().gather(gravityGradientDirection);
    const mfem::Vector displacementDirectionReduced = gravityContext.GetDisplacementMap().gather(displacementDirection);

    const auto layout                               = gravity_displacement_force_test_utils::make_layout(f);

    mean_field::operators::PreparedGravityDisplacementForceJacobianOperator adapter(layout, preparedOperator);

    mfem::BlockVector direction(layout.value_offsets());
    direction                                                                        = 0.0;

    direction.GetBlock(gravity_displacement_force_test_utils::densityValue)          = densityDirectionReduced;

    direction.GetBlock(gravity_displacement_force_test_utils::displacementValue)     = displacementDirectionReduced;

    direction.GetBlock(gravity_displacement_force_test_utils::gravityGradientValue)  = gravityGradientDirectionReduced;

    direction.GetBlock(gravity_displacement_force_test_utils::gravityPotentialValue) = 0.29;

    direction.GetBlock(gravity_displacement_force_test_utils::enthalpyValue)         = -0.37;

    direction.GetBlock(gravity_displacement_force_test_utils::barotropicConstantValue) = 0.43;

    mfem::Vector action;
    adapter.Mult(direction, action);

    mfem::Vector expectedDisplacementAction;

    preparedOperator.ApplyCompleteJacobianAction(
        densityDirectionReduced, displacementDirectionReduced, gravityGradientDirectionReduced,
        expectedDisplacementAction
    );

    const mfem::Vector actualDisplacementAction = gravity_displacement_force_test_utils::copy_residual_block(
        action, layout, gravity_displacement_force_test_utils::displacementResidual
    );

    CHECK(
        gravity_displacement_force_test_utils::relative_difference(
            actualDisplacementAction, expectedDisplacementAction, f.mesh->GetComm()
        ) < 2.0e-12
    );

    const std::array<mfem::Vector, 5> zeroRows{
        gravity_displacement_force_test_utils::copy_residual_block(
            action, layout, gravity_displacement_force_test_utils::gravityGradientResidual
        ),
        gravity_displacement_force_test_utils::copy_residual_block(
            action, layout, gravity_displacement_force_test_utils::gravityPotentialResidual
        ),
        gravity_displacement_force_test_utils::copy_residual_block(
            action, layout, gravity_displacement_force_test_utils::densityResidual
        ),
        gravity_displacement_force_test_utils::copy_residual_block(
            action, layout, gravity_displacement_force_test_utils::enthalpyResidual
        ),
        gravity_displacement_force_test_utils::copy_residual_block(
            action, layout, gravity_displacement_force_test_utils::massResidual
        )
    };

    for (const mfem::Vector &row : zeroRows) {
        CHECK(gravity_prepared_test_utils::global_norm(row, f.mesh->GetComm()) == 0.0);
    }
}