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

#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>
#include <mpi.h>

import mean_field;
import test_helpers;

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

    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

    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::field::FieldDofMap make_enthalpy_map(const mean_field::fem::FEM &f) {
        return mean_field::field::make_field_dof_map<mean_field::field::Enthalpy, DomainSchema>(*f.enthalpyFes);
    }

    [[nodiscard]] mean_field::operators::DisplacementResidualLayout 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 = make_enthalpy_map(f);

        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 field(f.densityFes.get());

        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.84 + 0.06 * std::sin(0.73 * position(0) + phase) + 0.04 * std::cos(0.61 * position(1) - phase) +
                   0.025 * position(2) * position(2);
        });

        field.ProjectCoefficient(coefficient);

        mfem::Vector trueDofs;
        field.GetTrueDofs(trueDofs);
        return trueDofs;
    }

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

        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.17 * std::sin(0.91 * position(0) + phase) - 0.11 * std::cos(0.79 * position(1) - phase) +
                   0.07 * position(2);
        });

        field.ProjectCoefficient(coefficient);

        mfem::Vector trueDofs;
        field.GetTrueDofs(trueDofs);
        return trueDofs;
    }

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

        auto function = [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 coefficient(3, function);
        field.ProjectCoefficient(coefficient);

        mfem::Vector trueDofs;
        field.GetTrueDofs(trueDofs);
        return trueDofs;
    }

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

        auto function = [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 coefficient(3, function);
        field.ProjectCoefficient(coefficient);

        mfem::Vector trueDofs;
        field.GetTrueDofs(trueDofs);
        return trueDofs;
    }

    [[nodiscard]] mfem::Vector make_positive_enthalpy(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        /*
         * Build a full H1 test state first, then return exactly the
         * solver-facing supported FieldDof coordinates consumed by the
         * migrated pressure-force operator.
         *
         * Keeping the public helper name unchanged means every existing
         * displacement-composer test automatically migrates to the new
         * pressure contract without inventing a parallel "*_full" helper.
         */
        mfem::Vector fullEnthalpy(f.enthalpyFes->GetTrueVSize());

        for (int index = 0; index < fullEnthalpy.Size(); ++index) {
            const double coordinate = static_cast<double>(index + 1);

            fullEnthalpy(index) =
                0.93 + 0.09 * std::sin(0.23 * coordinate + phase) + 0.04 * std::cos(0.17 * coordinate - 0.5 * phase);
        }

        return make_enthalpy_map(f).gather(fullEnthalpy);
    }

    [[nodiscard]] mfem::Vector make_enthalpy_direction(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::Vector fullDirection(f.enthalpyFes->GetTrueVSize());

        for (int index = 0; index < fullDirection.Size(); ++index) {
            const double coordinate = static_cast<double>(index + 1);

            fullDirection(index) =
                0.27 * std::sin(0.19 * coordinate + phase) + 0.14 * std::cos(0.13 * coordinate - 0.5 * phase);
        }

        return make_enthalpy_map(f).gather(fullDirection);
    }

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

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

        direction -= second;
        return direction;
    }

    [[nodiscard]] mean_field::physics::RigidRotation make_rotation(const double scale = 1.0) {
        mfem::Vector angularVelocity(3);
        angularVelocity(0) = scale * 0.17;
        angularVelocity(1) = scale * -0.09;
        angularVelocity(2) = scale * 0.62;

        mfem::Vector center(3);
        center(0) = 0.04;
        center(1) = -0.03;
        center(2) = 0.02;

        return mean_field::physics::RigidRotation(angularVelocity, center);
    }

    [[nodiscard]] mean_field::operators::DisplacementResidualDependencies make_dependencies() {
        return {
            .discretization  = {.identity = 401, .revision = 3},
            .density         = {.identity = 409, .revision = 5},
            .displacement    = {.identity = 419, .revision = 7},
            .gravityGradient = {.identity = 421, .revision = 11},
            .enthalpy        = {.identity = 431, .revision = 13},
            .rotation        = {.identity = 433, .revision = 17}
        };
    }

    [[nodiscard]] mean_field::operators::context::gravity_field::GravityFieldRevisions make_gravity_revisions(
        const mean_field::operators::DisplacementResidualDependencies &dependencies,
        const std::uint64_t potentialRevision
    ) {
        return {
            .discretization    = {.value = dependencies.discretization.revision},
            .displacement      = {.value = dependencies.displacement.revision},
            .density           = {.value = dependencies.density.revision},
            .gravity_gradient  = {.value = dependencies.gravityGradient.revision},
            .gravity_potential = {.value = potentialRevision}
        };
    }

    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::DisplacementResidualDependencies &dependencies,
        const std::uint64_t potentialRevision
    ) {
        context.Prepare(
            {.density           = context.GetDensityMap().gather(density),
             .displacement      = context.GetDisplacementMap().gather(displacement),
             .gravity_gradient  = context.GetGravityGradientMap().gather(gravityGradient),
             .gravity_potential = context.GetGravityPotentialMap().gather(gravityPotential)},
            make_gravity_revisions(dependencies, potentialRevision)
        );
    }

    [[nodiscard]] double relative_difference(
        const mfem::Vector &left,
        const mfem::Vector &right,
        const MPI_Comm communicator
    ) {
        REQUIRE(left.Size() == right.Size());

        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
    explicit_residual_sum(const mean_field::operators::PreparedDisplacementResidualOperator &preparedOperator) {
        mfem::Vector pressure;
        mfem::Vector gravity;
        mfem::Vector rotation;

        preparedOperator.GetPressureOperator().BuildResidual(pressure);
        preparedOperator.GetGravityOperator().BuildResidual(gravity);
        preparedOperator.GetRotationalOperator().BuildResidual(rotation);

        pressure += gravity;
        pressure += rotation;
        return pressure;
    }

    template <int index>
    [[nodiscard]] mfem::Vector copy_residual_block(
        const mfem::Vector &action,
        const mean_field::operators::DisplacementResidualLayout &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 prepared_displacement_residual_test_utils

TEST_CASE(
    "Prepared Displacement Residual Equals The Three Prepared Contributors",
    tags::barotrope_prepared
) {
    using Operator = mean_field::operators::PreparedDisplacementResidualOperator;

    STATIC_REQUIRE_FALSE(std::is_copy_constructible_v<Operator>);
    STATIC_REQUIRE_FALSE(std::is_copy_assignable_v<Operator>);
    STATIC_REQUIRE_FALSE(std::is_move_constructible_v<Operator>);
    STATIC_REQUIRE_FALSE(std::is_move_assignable_v<Operator>);

    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 auto enthalpyMap             = prepared_displacement_residual_test_utils::make_enthalpy_map(f);

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

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

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

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

    const mfem::Vector enthalpy = prepared_displacement_residual_test_utils::make_positive_enthalpy(f, 0.53);

    REQUIRE(enthalpyMap.reduced_size() < enthalpyMap.full_size());
    REQUIRE(enthalpy.Size() == enthalpyMap.reduced_size());

    const auto dependencies = prepared_displacement_residual_test_utils::make_dependencies();

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

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, 19
    );

    const mean_field::eos::Polytrope equationOfState(3.0, 0.25);
    const mean_field::physics::RigidRotation rotation = prepared_displacement_residual_test_utils::make_rotation(0.83);

    Operator preparedOperator(f, *f.domainMapperStateless, equationOfState, gravityContext);

    REQUIRE(preparedOperator.GetPressureOperator().GetEnthalpySize() == enthalpy.Size());

    const auto initialReport = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    REQUIRE(initialReport.pressure.DidAnyWork());
    REQUIRE(initialReport.gravity.DidAnyWork());
    REQUIRE(initialReport.rotation.DidAnyWork());
    REQUIRE(initialReport.assembledResidual);
    REQUIRE(preparedOperator.IsPrepared());

    CHECK(&preparedOperator.GetFEM() == &f);
    CHECK(&preparedOperator.GetGravityContext() == &gravityContext);
    CHECK(&preparedOperator.GetGravityOperator().GetGravityContext() == &gravityContext);

    mfem::Vector compositeResidual;
    preparedOperator.BuildResidual(compositeResidual);

    const mfem::Vector explicitResidual =
        prepared_displacement_residual_test_utils::explicit_residual_sum(preparedOperator);

    const double compositionError = prepared_displacement_residual_test_utils::relative_difference(
        compositeResidual, explicitResidual, f.mesh->GetComm()
    );

    INFO("Prepared residual composition error = " << compositionError);
    CHECK(compositionError < 2.0e-15);

    const std::uint64_t preparationCount = preparedOperator.GetResidualPreparationCount();

    const auto repeatedReport            = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    CHECK_FALSE(repeatedReport.DidAnyWork());
    CHECK_FALSE(repeatedReport.assembledResidual);
    CHECK(preparedOperator.GetResidualPreparationCount() == preparationCount);
    CHECK(preparedOperator.IsPrepared());
}

TEST_CASE(
    "Prepared Displacement Residual Preserves Pressure Rejection Details",
    tags::barotrope_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());

    const mfem::Vector density         = prepared_displacement_residual_test_utils::make_density(f, 0.31);
    const mfem::Vector displacement    = gravity_prepared_test_utils::make_displacement(f, 0.67);
    const mfem::Vector gravityGradient = prepared_displacement_residual_test_utils::make_gravity_gradient(f, 0.47);
    mfem::Vector gravityPotential(f.gravityPotentialFes->GetTrueVSize());
    gravityPotential  = 0.0;

    auto dependencies = prepared_displacement_residual_test_utils::make_dependencies();
    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );
    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, 19
    );

    const mean_field::eos::Polytrope equationOfState(3.0, 0.25);
    const mean_field::physics::RigidRotation rotation = prepared_displacement_residual_test_utils::make_rotation(0.83);
    mean_field::operators::PreparedDisplacementResidualOperator preparedOperator(
        f, *f.domainMapperStateless, equationOfState, gravityContext
    );

    mfem::Vector enthalpy(prepared_displacement_residual_test_utils::make_enthalpy_map(f).reduced_size());
    enthalpy            = std::numeric_limits<double>::max();

    const auto rejected = preparedOperator.TryPrepare({.enthalpy = enthalpy}, dependencies, rotation);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().source == mean_field::operators::DisplacementResidualPreparationRejectionSource::pressure);
    CHECK(
        rejected.error().reason ==
        mean_field::operators::DisplacementResidualPreparationRejectionReason::equation_of_state
    );
    CHECK(rejected.error().equationOfStateCode == mean_field::eos::EvaluationErrorCode::nonfinite_result);
    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK_THROWS_AS(
        preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation), mean_field::eos::EvaluationError
    );

    enthalpy = 1.0;
    ++dependencies.enthalpy.revision;
    const auto accepted = preparedOperator.TryPrepare({.enthalpy = enthalpy}, dependencies, rotation);
    REQUIRE(accepted.has_value());
    CHECK(preparedOperator.IsPrepared());
}

TEST_CASE(
    "Prepared Displacement Residual Selectively Orchestrates Its Children",
    tags::barotrope_context_integration
) {
    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         = prepared_displacement_residual_test_utils::make_density(f, 0.29);

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

    mfem::Vector gravityGradient = prepared_displacement_residual_test_utils::make_gravity_gradient(f, 0.43);

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

    mfem::Vector enthalpy           = prepared_displacement_residual_test_utils::make_positive_enthalpy(f, 0.51);

    auto dependencies               = prepared_displacement_residual_test_utils::make_dependencies();

    std::uint64_t potentialRevision = 19;

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

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, potentialRevision
    );

    const mean_field::eos::Polytrope equationOfState(3.0, 0.25);

    mean_field::physics::RigidRotation rotation = prepared_displacement_residual_test_utils::make_rotation(0.79);

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

    preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    gravityPotential = 0.17;
    ++potentialRevision;

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, potentialRevision
    );

    const auto potentialReport = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    CHECK_FALSE(potentialReport.DidAnyWork());
    CHECK_FALSE(potentialReport.assembledResidual);

    enthalpy = prepared_displacement_residual_test_utils::make_positive_enthalpy(f, 0.83);
    ++dependencies.enthalpy.revision;

    const auto enthalpyReport = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    CHECK(enthalpyReport.pressure.DidAnyWork());
    CHECK_FALSE(enthalpyReport.gravity.DidAnyWork());
    CHECK_FALSE(enthalpyReport.rotation.DidAnyWork());
    CHECK(enthalpyReport.assembledResidual);

    gravityGradient = prepared_displacement_residual_test_utils::make_gravity_gradient(f, 0.91);
    ++dependencies.gravityGradient.revision;

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, potentialRevision
    );

    const auto gravityReport = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    CHECK_FALSE(gravityReport.pressure.DidAnyWork());
    CHECK(gravityReport.gravity.DidAnyWork());
    CHECK_FALSE(gravityReport.rotation.DidAnyWork());
    CHECK(gravityReport.assembledResidual);

    density = prepared_displacement_residual_test_utils::make_density(f, 1.07);
    ++dependencies.density.revision;

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, potentialRevision
    );

    const auto densityReport = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    CHECK_FALSE(densityReport.pressure.DidAnyWork());
    CHECK(densityReport.gravity.DidAnyWork());
    CHECK(densityReport.rotation.DidAnyWork());
    CHECK(densityReport.assembledResidual);

    rotation = prepared_displacement_residual_test_utils::make_rotation(1.13);
    ++dependencies.rotation.revision;

    const auto rotationReport = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    CHECK_FALSE(rotationReport.pressure.DidAnyWork());
    CHECK_FALSE(rotationReport.gravity.DidAnyWork());
    CHECK(rotationReport.rotation.DidAnyWork());
    CHECK(rotationReport.assembledResidual);

    displacement = gravity_prepared_test_utils::make_displacement(f, 0.89);
    ++dependencies.displacement.revision;

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, potentialRevision
    );

    const auto displacementReport = preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

    CHECK(displacementReport.pressure.DidAnyWork());
    CHECK(displacementReport.gravity.DidAnyWork());
    CHECK(displacementReport.rotation.DidAnyWork());
    CHECK(displacementReport.assembledResidual);
}

TEST_CASE(
    "Prepared Displacement Residual Jacobian Equals The Contributor Sums",
    tags::barotrope_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         = prepared_displacement_residual_test_utils::make_density(f, 0.37);

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

    const mfem::Vector gravityGradient = prepared_displacement_residual_test_utils::make_gravity_gradient(f, 0.59);

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

    const mfem::Vector enthalpy         = prepared_displacement_residual_test_utils::make_positive_enthalpy(f, 0.61);

    const mfem::Vector densityDirection = prepared_displacement_residual_test_utils::make_density_direction(f, 0.71);

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

    const mfem::Vector gravityDirection =
        prepared_displacement_residual_test_utils::make_gravity_gradient_direction(f, 0.83);

    const mfem::Vector enthalpyDirection = prepared_displacement_residual_test_utils::make_enthalpy_direction(f, 0.97);

    const auto dependencies              = prepared_displacement_residual_test_utils::make_dependencies();

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

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, 19
    );

    const mean_field::eos::Polytrope equationOfState(3.0, 0.25);
    const mean_field::physics::RigidRotation rotation = prepared_displacement_residual_test_utils::make_rotation(0.91);

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

    preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

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

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

    preparedOperator.ApplyDensityJacobianAction(densityDirectionReduced, densityAction);

    preparedOperator.ApplyDisplacementJacobianAction(displacementDirectionReduced, displacementAction);

    preparedOperator.ApplyGravityGradientJacobianAction(gravityDirectionReduced, gravityAction);

    preparedOperator.ApplyEnthalpyJacobianAction(enthalpyDirection, enthalpyAction);

    preparedOperator.ApplyCompleteJacobianAction(
        densityDirectionReduced, displacementDirectionReduced, gravityDirectionReduced, enthalpyDirection,
        completeAction
    );

    mfem::Vector expectedDensity;
    mfem::Vector expectedRotationDensity;

    preparedOperator.GetGravityOperator().ApplyDensityJacobianAction(densityDirectionReduced, expectedDensity);

    preparedOperator.GetRotationalOperator().ApplyDensityJacobianAction(
        densityDirectionReduced, expectedRotationDensity
    );

    expectedDensity += expectedRotationDensity;

    mfem::Vector expectedDisplacement;
    mfem::Vector expectedGravityDisplacement;
    mfem::Vector expectedRotationDisplacement;

    preparedOperator.GetPressureOperator().ApplyDisplacementJacobianAction(
        displacementDirectionReduced, expectedDisplacement
    );

    preparedOperator.GetGravityOperator().ApplyDisplacementJacobianAction(
        displacementDirectionReduced, expectedGravityDisplacement
    );

    preparedOperator.GetRotationalOperator().ApplyDisplacementJacobianAction(
        displacementDirectionReduced, expectedRotationDisplacement
    );

    expectedDisplacement += expectedGravityDisplacement;
    expectedDisplacement += expectedRotationDisplacement;

    mfem::Vector expectedGravity;
    preparedOperator.GetGravityOperator().ApplyGravityGradientJacobianAction(gravityDirectionReduced, expectedGravity);

    mfem::Vector expectedEnthalpy;
    preparedOperator.GetPressureOperator().ApplyEnthalpyJacobianAction(enthalpyDirection, expectedEnthalpy);

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

    const MPI_Comm communicator = f.mesh->GetComm();

    CHECK(
        prepared_displacement_residual_test_utils::relative_difference(densityAction, expectedDensity, communicator) <
        2.0e-15
    );

    CHECK(
        prepared_displacement_residual_test_utils::relative_difference(
            displacementAction, expectedDisplacement, communicator
        ) < 2.0e-15
    );

    CHECK(
        prepared_displacement_residual_test_utils::relative_difference(gravityAction, expectedGravity, communicator) <
        2.0e-15
    );

    CHECK(
        prepared_displacement_residual_test_utils::relative_difference(enthalpyAction, expectedEnthalpy, communicator) <
        2.0e-15
    );

    CHECK(
        prepared_displacement_residual_test_utils::relative_difference(completeAction, summedColumns, communicator) <
        2.0e-15
    );
}

TEST_CASE(
    "Prepared Displacement Residual Jacobian Matches A Simultaneous "
    "Centered Difference On Deformed Geometry",
    tags::barotrope_prepared_jacobian_accuracy &tags::geometry
) {
    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 baseDensity      = prepared_displacement_residual_test_utils::make_density(f, 0.41);

    const mfem::Vector baseDisplacement = gravity_prepared_test_utils::make_displacement(f, 0.79);

    const mfem::Vector baseGravity      = prepared_displacement_residual_test_utils::make_gravity_gradient(f, 0.63);

    const mfem::Vector baseEnthalpy     = prepared_displacement_residual_test_utils::make_positive_enthalpy(f, 0.67);

    const mfem::Vector densityDirection = prepared_displacement_residual_test_utils::make_density_direction(f, 0.73);

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

    const mfem::Vector gravityDirection =
        prepared_displacement_residual_test_utils::make_gravity_gradient_direction(f, 0.89);

    const mfem::Vector enthalpyDirection = prepared_displacement_residual_test_utils::make_enthalpy_direction(f, 1.01);

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

    auto dependencies = prepared_displacement_residual_test_utils::make_dependencies();

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

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, baseDensity, baseDisplacement, baseGravity, gravityPotential, dependencies, 19
    );

    const mean_field::eos::Polytrope equationOfState(3.0, 0.25);
    const mean_field::physics::RigidRotation rotation = prepared_displacement_residual_test_utils::make_rotation(0.87);

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

    preparedOperator.Prepare({.enthalpy = baseEnthalpy}, dependencies, rotation);

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

    mfem::Vector jacobianAction;

    preparedOperator.ApplyCompleteJacobianAction(
        densityDirectionReduced, displacementDirectionReduced, gravityDirectionReduced, enthalpyDirection,
        jacobianAction
    );

    constexpr double step = 1.0e-5;

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

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

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

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

    mfem::Vector plusGravity(baseGravity);
    plusGravity.Add(step, gravityDirection);

    mfem::Vector minusGravity(baseGravity);
    minusGravity.Add(-step, gravityDirection);

    mfem::Vector plusEnthalpy(baseEnthalpy);
    plusEnthalpy.Add(step, enthalpyDirection);

    mfem::Vector minusEnthalpy(baseEnthalpy);
    minusEnthalpy.Add(-step, enthalpyDirection);

    ++dependencies.density.revision;
    ++dependencies.displacement.revision;
    ++dependencies.gravityGradient.revision;
    ++dependencies.enthalpy.revision;

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, plusDensity, plusDisplacement, plusGravity, gravityPotential, dependencies, 19
    );

    preparedOperator.Prepare({.enthalpy = plusEnthalpy}, dependencies, rotation);

    mfem::Vector plusResidual;
    preparedOperator.BuildResidual(plusResidual);

    ++dependencies.density.revision;
    ++dependencies.displacement.revision;
    ++dependencies.gravityGradient.revision;
    ++dependencies.enthalpy.revision;

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, minusDensity, minusDisplacement, minusGravity, gravityPotential, dependencies, 19
    );

    preparedOperator.Prepare({.enthalpy = minusEnthalpy}, dependencies, rotation);

    mfem::Vector minusResidual;
    preparedOperator.BuildResidual(minusResidual);

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

    const double centeredDifferenceError =
        prepared_displacement_residual_test_utils::relative_difference(jacobianAction, plusResidual, f.mesh->GetComm());

    INFO("Composite simultaneous centered-difference error = " << centeredDifferenceError);

    CHECK(centeredDifferenceError < 8.0e-8);
}

TEST_CASE(
    "Prepared Displacement Residual MFEM Adapter Routes Only R-d",
    tags::barotrope_prepared_jacobian_unit
) {
    using JacobianForm             = mean_field::utils::blocks::barotropic_equilibrium_jacobian_form;

    using DisplacementResidualType = mean_field::utils::blocks::displacement::geometry::residual;

    STATIC_REQUIRE(
        mean_field::utils::blocks::has_jacobian_coupling_v<
            DisplacementResidualType, mean_field::utils::blocks::density::mass::value, JacobianForm>
    );

    STATIC_REQUIRE(
        mean_field::utils::blocks::has_jacobian_coupling_v<
            DisplacementResidualType, mean_field::utils::blocks::displacement::geometry::value, JacobianForm>
    );

    STATIC_REQUIRE(
        mean_field::utils::blocks::has_jacobian_coupling_v<
            DisplacementResidualType, mean_field::utils::blocks::gravity::gradient::value, JacobianForm>
    );

    STATIC_REQUIRE(
        mean_field::utils::blocks::has_jacobian_coupling_v<
            DisplacementResidualType, mean_field::utils::blocks::enthalpy::specific::value, JacobianForm>
    );

    STATIC_REQUIRE_FALSE(
        mean_field::utils::blocks::has_jacobian_coupling_v<
            DisplacementResidualType, mean_field::utils::blocks::gravity::poisson::value, JacobianForm>
    );

    STATIC_REQUIRE_FALSE(
        mean_field::utils::blocks::has_jacobian_coupling_v<
            DisplacementResidualType, mean_field::utils::blocks::barotropic_constant::mass_normalization::value,
            JacobianForm>
    );

    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         = prepared_displacement_residual_test_utils::make_density(f, 0.43);

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

    const mfem::Vector gravityGradient = prepared_displacement_residual_test_utils::make_gravity_gradient(f, 0.57);

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

    const mfem::Vector enthalpy         = prepared_displacement_residual_test_utils::make_positive_enthalpy(f, 0.69);

    const mfem::Vector densityDirection = prepared_displacement_residual_test_utils::make_density_direction(f, 0.77);

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

    const mfem::Vector gravityDirection =
        prepared_displacement_residual_test_utils::make_gravity_gradient_direction(f, 0.93);

    const mfem::Vector enthalpyDirection = prepared_displacement_residual_test_utils::make_enthalpy_direction(f, 1.03);

    const auto dependencies              = prepared_displacement_residual_test_utils::make_dependencies();

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

    prepared_displacement_residual_test_utils::prepare_gravity_context(
        gravityContext, density, displacement, gravityGradient, gravityPotential, dependencies, 19
    );

    const mean_field::eos::Polytrope equationOfState(3.0, 0.25);
    const mean_field::physics::RigidRotation rotation = prepared_displacement_residual_test_utils::make_rotation(0.95);

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

    preparedOperator.Prepare({.enthalpy = enthalpy}, dependencies, rotation);

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

    const mean_field::operators::DisplacementResidualLayout layout =
        prepared_displacement_residual_test_utils::make_layout(f);

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

    CHECK(adapter.Width() == layout.value_offsets().Last());
    CHECK(adapter.Height() == layout.residual_offsets().Last());
    CHECK(
        layout.size(prepared_displacement_residual_test_utils::enthalpyValue) ==
        preparedOperator.GetPressureOperator().GetEnthalpySize()
    );

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

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

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

    direction.GetBlock(prepared_displacement_residual_test_utils::gravityGradientValue)  = gravityDirectionReduced;

    direction.GetBlock(prepared_displacement_residual_test_utils::enthalpyValue)         = enthalpyDirection;

    direction.GetBlock(prepared_displacement_residual_test_utils::gravityPotentialValue) = 0.59;

    direction.GetBlock(prepared_displacement_residual_test_utils::barotropicConstantValue) = -0.73;

    mfem::Vector expectedDisplacementAction;

    preparedOperator.ApplyCompleteJacobianAction(
        densityDirectionReduced, displacementDirectionReduced, gravityDirectionReduced, enthalpyDirection,
        expectedDisplacementAction
    );

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

    const mfem::Vector routedDisplacement = prepared_displacement_residual_test_utils::copy_residual_block(
        action, layout, prepared_displacement_residual_test_utils::displacementResidual
    );

    CHECK(
        prepared_displacement_residual_test_utils::relative_difference(
            routedDisplacement, expectedDisplacementAction, f.mesh->GetComm()
        ) < 2.0e-15
    );

    const mfem::Vector gravityGradientBlock = prepared_displacement_residual_test_utils::copy_residual_block(
        action, layout, prepared_displacement_residual_test_utils::gravityGradientResidual
    );

    const mfem::Vector gravityPotentialBlock = prepared_displacement_residual_test_utils::copy_residual_block(
        action, layout, prepared_displacement_residual_test_utils::gravityPotentialResidual
    );

    const mfem::Vector densityBlock = prepared_displacement_residual_test_utils::copy_residual_block(
        action, layout, prepared_displacement_residual_test_utils::densityResidual
    );

    const mfem::Vector enthalpyBlock = prepared_displacement_residual_test_utils::copy_residual_block(
        action, layout, prepared_displacement_residual_test_utils::enthalpyResidual
    );

    const mfem::Vector massBlock = prepared_displacement_residual_test_utils::copy_residual_block(
        action, layout, prepared_displacement_residual_test_utils::massResidual
    );

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

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

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

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

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