module;

#include <cmath>
#include <expected>
#include <optional>
#include <stdexcept>
#include <utility>

#include <mfem.hpp>
#include <mpi.h>

module mean_field;

import :operators.kernels.gravity_displacement_force;
import :operators.prepared_gravity_displacement_force;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using Rejection    = mean_field::operators::kernels::GravityDisplacementForceRejection;
    using Reason       = mean_field::operators::kernels::GravityDisplacementForceRejectionReason;

    [[nodiscard]] bool relevant_revisions_match(
        const mean_field::operators::context::gravity_field::GravityFieldRevisions &left,
        const mean_field::operators::context::gravity_field::GravityFieldRevisions &right
    ) noexcept {
        return left.discretization == right.discretization && left.displacement == right.displacement &&
               left.density == right.density && left.gravity_gradient == right.gravity_gradient;
    }

    [[nodiscard]] bool is_vacuum_attribute(const int attribute) {
        return DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Vacuum>(attribute);
    }

    [[nodiscard]] Rejection mapping_rejection(const mean_field::mapping::MappingStatus status) {
        MFEM_VERIFY(
            status != mean_field::mapping::MappingStatus::invalid_dimension,
            "Prepared gravity force mapping reported an invariant dimension mismatch."
        );
        return {.reason = Reason::invalid_mapping, .mappingStatus = status};
    }

    [[nodiscard]] Rejection non_finite_rejection() noexcept {
        return {.reason = Reason::non_finite_arithmetic};
    }

    [[nodiscard]] int encode_rejection(const std::optional<Rejection> &rejection) noexcept {
        if (!rejection.has_value()) {
            return 0;
        }
        if (rejection->reason == Reason::non_finite_arithmetic) {
            return 256;
        }
        return static_cast<int>(rejection->mappingStatus) + 1;
    }

    [[nodiscard]] Rejection decode_rejection(const int encoded) {
        if (encoded >= 256) {
            return non_finite_rejection();
        }
        return mapping_rejection(static_cast<mean_field::mapping::MappingStatus>(encoded - 1));
    }

    [[nodiscard]] std::expected<
        void,
        Rejection>
    synchronize_rejection(
        const std::optional<Rejection> &localRejection,
        const MPI_Comm communicator
    ) {
        const int localEncoded = encode_rejection(localRejection);
        int globalEncoded      = 0;
        if (MPI_Allreduce(&localEncoded, &globalEncoded, 1, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS) {
            throw std::runtime_error("Could not synchronize prepared gravity-force candidate validity.");
        }
        if (globalEncoded != 0) {
            return std::unexpected(decode_rejection(globalEncoded));
        }
        return {};
    }

    [[nodiscard]] bool vector_is_finite(const mfem::Vector &vector) noexcept {
        for (int index = 0; index < vector.Size(); ++index) {
            if (!std::isfinite(vector(index))) {
                return false;
            }
        }
        return true;
    }

    [[noreturn]] void throw_rejection(const Rejection &rejection) {
        if (rejection.reason == Reason::non_finite_arithmetic) {
            throw std::domain_error("Prepared gravity force produced non-finite arithmetic.");
        }
        throw std::domain_error("Prepared gravity force encountered an invalid mapped domain.");
    }

    void true_to_local(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &trueVector,
        mfem::Vector &localVector
    ) {
        localVector.SetSize(finiteElementSpace.GetVSize());
        const mfem::Operator *prolongation = finiteElementSpace.GetProlongationMatrix();
        if (prolongation != nullptr) {
            prolongation->Mult(trueVector, localVector);
        } else {
            localVector = trueVector;
        }
    }

    void local_to_true(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &localVector,
        mfem::Vector &trueVector
    ) {
        trueVector.SetSize(finiteElementSpace.GetTrueVSize());
        trueVector                         = 0.0;
        const mfem::Operator *prolongation = finiteElementSpace.GetProlongationMatrix();
        if (prolongation != nullptr) {
            prolongation->MultTranspose(localVector, trueVector);
        } else {
            trueVector = localVector;
        }
    }

    [[nodiscard]] int vector_dof_index(
        const mfem::Ordering::Type ordering,
        const int scalarDof,
        const int component,
        const int scalarDofCount,
        const int dimension
    ) {
        if (ordering == mfem::Ordering::byNODES) {
            return scalarDof + component * scalarDofCount;
        }
        MFEM_VERIFY(ordering == mfem::Ordering::byVDIM, "Unsupported displacement ordering.");
        return scalarDof * dimension + component;
    }

    [[nodiscard]] const mfem::IntegrationRule &get_gravity_force_rule(
        const mean_field::fem::FEM &f,
        const mfem::ElementTransformation &transformation
    ) {
        using DisplacementField = mean_field::field::Field<mean_field::field::Displacement>;
        const mean_field::quadrature::Query query =
            DisplacementField::make_query<mean_field::field::Displacement::Form::GravityForce>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
                mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
            );
        const mean_field::quadrature::MfemRule rule = f.quadratureFactory->get(query, transformation.GetGeometryType());
        MFEM_VERIFY(
            rule.integration_rule != nullptr,
            "The quadrature policy did not return a gravity-displacement-force integration rule."
        );
        return *rule.integration_rule;
    }
} // namespace

namespace mean_field::operators {
    PreparedGravityDisplacementForceOperator::PreparedGravityDisplacementForceOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const context::gravity_field::GravityFieldLinearizationContext &gravityContext
    )
        : m_fem(f),
          m_domainMapper(domainMapper),
          m_gravityContext(gravityContext) {
        MFEM_VERIFY(m_fem.mesh != nullptr, "PreparedGravityDisplacementForceOperator requires a mesh.");

        MFEM_VERIFY(
            m_fem.densityFes != nullptr && m_fem.gravityFluxFes != nullptr && m_fem.displacementFes != nullptr,
            "PreparedGravityDisplacementForceOperator requires density, "
            "gravity-gradient, and displacement finite-element spaces."
        );

        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_fem.mesh->Dimension(),
            "PreparedGravityDisplacementForceOperator received a mapper "
            "with the wrong dimension."
        );
    }

    std::expected<
        void,
        kernels::GravityDisplacementForceRejection>
    PreparedGravityDisplacementForceOperator::TryPrepareElementData() {
        m_elements.clear();
        m_elements.reserve(m_fem.mesh->GetNE());

        mfem::Vector baseDensityLocal;
        mfem::Vector baseGravityGradientLocal;
        mfem::Vector baseDisplacementLocal;
        true_to_local(*m_fem.densityFes, m_gravityContext.GetDensityTrue(), baseDensityLocal);
        true_to_local(*m_fem.gravityFluxFes, m_gravityContext.GetGravityGradientTrue(), baseGravityGradientLocal);
        true_to_local(
            *m_fem.displacementFes, m_gravityContext.GetGeometryContext().GetDisplacementTrue(), baseDisplacementLocal
        );

        mapping::DomainMapper::Workspace workspace(m_domainMapper.GetDimension());
        mapping::VolumeMappingContext mappingContext;
        mfem::Array<int> compactificationDofs;
        mfem::Vector elementBaseDensity;
        mfem::Vector elementBaseGravityGradient;
        mfem::Vector elementBaseDisplacement;
        mfem::Vector elementCompactification;
        mfem::Vector densityShape;
        mfem::Vector baseGravityReferenceValue;
        mfem::DenseMatrix gravityGradientShape;

        const int dimension = m_domainMapper.GetDimension();

        for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(elementId);
            MFEM_VERIFY(transformation != nullptr, "Prepared gravity force received a null transformation.");
            if (is_vacuum_attribute(transformation->Attribute)) {
                continue;
            }

            m_elements.emplace_back();
            ElementPAData &data           = m_elements.back();
            data.elementId                = elementId;
            data.densityDofTransformation = m_fem.densityFes->GetElementDofs(elementId, data.densityDofs);
            data.gravityGradientDofTransformation =
                m_fem.gravityFluxFes->GetElementVDofs(elementId, data.gravityGradientDofs);
            data.displacementDofTransformation =
                m_fem.displacementFes->GetElementVDofs(elementId, data.displacementDofs);
            mfem::DofTransformation *compactificationDofTransformation =
                m_fem.compactificationFes->GetElementDofs(elementId, compactificationDofs);

            baseDensityLocal.GetSubVector(data.densityDofs, elementBaseDensity);
            baseGravityGradientLocal.GetSubVector(data.gravityGradientDofs, elementBaseGravityGradient);
            baseDisplacementLocal.GetSubVector(data.displacementDofs, elementBaseDisplacement);
            m_fem.compactificationCoordinate->GetSubVector(compactificationDofs, elementCompactification);
            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(elementBaseDensity);
            }
            if (data.gravityGradientDofTransformation != nullptr) {
                data.gravityGradientDofTransformation->InvTransformPrimal(elementBaseGravityGradient);
            }
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(elementBaseDisplacement);
            }
            if (compactificationDofTransformation != nullptr) {
                compactificationDofTransformation->InvTransformPrimal(elementCompactification);
            }

            const mfem::FiniteElement &densityElement          = *m_fem.densityFes->GetFE(elementId);
            const mfem::FiniteElement &gravityGradientElement  = *m_fem.gravityFluxFes->GetFE(elementId);
            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(elementId);
            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(elementId);
            data.integrationRule                               = &get_gravity_force_rule(m_fem, *transformation);

            const mapping::ElementDisplacementData displacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementBaseDisplacement);
            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, elementCompactification
            );
            const mapping::ElementMappingData mappingData{
                .displacement = displacementData, .compactification = compactificationData
            };

            const int quadraturePointCount = data.integrationRule->GetNPoints();
            data.mappingJacobians.SetSize(quadraturePointCount, dimension * dimension);
            data.inverseMeshJacobians.SetSize(quadraturePointCount, dimension * dimension);
            data.baseGravityReferenceValues.SetSize(quadraturePointCount, dimension);
            data.baseDensityValues.SetSize(quadraturePointCount);
            data.referenceWeights.SetSize(quadraturePointCount);
            densityShape.SetSize(densityElement.GetDof());
            gravityGradientShape.SetSize(gravityGradientElement.GetDof(), dimension);
            baseGravityReferenceValue.SetSize(dimension);

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
                const mapping::MappingStatus status            = m_domainMapper.EvaluateVolume(
                    mappingData, *transformation, integrationPoint, workspace, mappingContext
                );
                if (status != mapping::MappingStatus::valid) {
                    return std::unexpected(mapping_rejection(status));
                }
                MFEM_VERIFY(
                    !mappingContext.mapping.compactified,
                    "Prepared gravity force encountered compactification on a stellar element."
                );

                densityElement.CalcShape(integrationPoint, densityShape);
                gravityGradientElement.CalcVShape(*transformation, gravityGradientShape);
                gravityGradientShape.MultTranspose(elementBaseGravityGradient, baseGravityReferenceValue);
                data.baseDensityValues(quadraturePoint)      = elementBaseDensity * densityShape;
                data.referenceWeights(quadraturePoint)       = integrationPoint.weight * transformation->Weight();

                const mfem::DenseMatrix &inverseMeshJacobian = transformation->InverseJacobian();
                for (int row = 0; row < dimension; ++row) {
                    data.baseGravityReferenceValues(quadraturePoint, row) = baseGravityReferenceValue(row);
                    for (int column = 0; column < dimension; ++column) {
                        const int entry = row * dimension + column;
                        data.mappingJacobians(quadraturePoint, entry) =
                            mappingContext.mapping.mapping_jacobian(row, column);
                        data.inverseMeshJacobians(quadraturePoint, entry) = inverseMeshJacobian(row, column);
                    }
                }

                if (!std::isfinite(data.baseDensityValues(quadraturePoint)) ||
                    !std::isfinite(data.referenceWeights(quadraturePoint)) ||
                    !vector_is_finite(baseGravityReferenceValue)) {
                    return std::unexpected(non_finite_rejection());
                }
                for (int row = 0; row < dimension; ++row) {
                    for (int column = 0; column < dimension; ++column) {
                        const int entry = row * dimension + column;
                        if (!std::isfinite(data.mappingJacobians(quadraturePoint, entry)) ||
                            !std::isfinite(data.inverseMeshJacobians(quadraturePoint, entry))) {
                            return std::unexpected(non_finite_rejection());
                        }
                    }
                }
            }
        }
        return {};
    }

    PreparedGravityDisplacementForceReport PreparedGravityDisplacementForceOperator::Prepare() {
        auto result = TryPrepare();
        if (!result.has_value()) {
            throw_rejection(result.error());
        }
        return std::move(result).value();
    }

    std::expected<
        PreparedGravityDisplacementForceReport,
        kernels::GravityDisplacementForceRejection>
    PreparedGravityDisplacementForceOperator::TryPrepare() {
        MFEM_VERIFY(
            m_gravityContext.IsPrepared(), "PreparedGravityDisplacementForceOperator requires the shared "
                                           "gravity linearization context to be prepared first."
        );

        const context::gravity_field::GravityFieldRevisions &requestedRevisions = m_gravityContext.GetRevisions();

        if (m_isPrepared && relevant_revisions_match(requestedRevisions, m_preparedRevisions)) {
            return {};
        }

        m_isPrepared             = false;

        /*
         * Build the reusable element plan before assembling the residual.
         * This pass stops at the first invalid mapped quadrature point, so a
         * rejected line-search candidate need not traverse the full stateless
         * residual kernel.  Synchronize before proceeding so every rank takes
         * the same branch.
         */
        const auto elementResult = TryPrepareElementData();
        const std::optional<Rejection> localElementRejection =
            elementResult.has_value() ? std::optional<Rejection>{} : std::optional<Rejection>{elementResult.error()};
        auto synchronizedElement = synchronize_rejection(localElementRejection, m_fem.mesh->GetComm());
        if (!synchronizedElement.has_value()) {
            return std::unexpected(synchronizedElement.error());
        }

        auto residualResult = kernels::try_apply_gravity_displacement_force_residual(
            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityContext.GetGravityGradientTrue(),
            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        if (!residualResult.has_value()) {
            return std::unexpected(residualResult.error());
        }
        m_cachedResidual.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, m_cachedResidual);
        std::optional<Rejection> localRejection;
        if (!vector_is_finite(m_cachedResidual)) {
            localRejection = non_finite_rejection();
        }
        auto synchronized = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
        if (!synchronized.has_value()) {
            return std::unexpected(synchronized.error());
        }
        m_preparedRevisions = requestedRevisions;
        ++m_residualPreparationCount;
        m_isPrepared = true;

        return PreparedGravityDisplacementForceReport{.preparedResidual = true};
    }

    void PreparedGravityDisplacementForceOperator::BuildResidual(mfem::Vector &residual) const {
        VerifyPrepared();
        residual = m_cachedResidual;
        ++m_residualApplicationCount;
    }

    void PreparedGravityDisplacementForceOperator::ApplyDensityJacobianAction(
        const mfem::Vector &densityVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);

        kernels::apply_gravity_displacement_force_density_action(
            m_fem, m_domainMapper, m_densityVariationTrue, m_gravityContext.GetGravityGradientTrue(),
            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_densityJacobianStatistics.applications;
    }

    void PreparedGravityDisplacementForceOperator::ApplyGravityGradientJacobianAction(
        const mfem::Vector &gravityGradientVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_gravityGradientVariationTrue.SetSize(m_gravityContext.GetGravityGradientMap().full_size());
        m_gravityContext.GetGravityGradientMap().scatter(gravityGradientVariation, m_gravityGradientVariationTrue);

        kernels::apply_gravity_displacement_force_gradient_action(
            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityGradientVariationTrue,
            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_gravityGradientJacobianStatistics.applications;
    }

    void PreparedGravityDisplacementForceOperator::ApplyDisplacementJacobianAction(
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        kernels::apply_gravity_displacement_force_displacement_action(
            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityContext.GetGravityGradientTrue(),
            m_displacementVariationTrue, m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_displacementJacobianStatistics.applications;
    }

    void PreparedGravityDisplacementForceOperator::ApplyPreparedCompleteJacobianActionTrue(
        const mfem::Vector &densityVariationTrue,
        const mfem::Vector &displacementVariationTrue,
        const mfem::Vector &gravityGradientVariationTrue,
        mfem::Vector &actionTrue
    ) const {
        true_to_local(*m_fem.densityFes, densityVariationTrue, m_densityVariationLocal);
        true_to_local(*m_fem.gravityFluxFes, gravityGradientVariationTrue, m_gravityGradientVariationLocal);
        true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacementVariationLocal);
        m_localAction.SetSize(m_fem.displacementFes->GetVSize());
        m_localAction                       = 0.0;

        const int dimension                 = m_domainMapper.GetDimension();
        const mfem::Ordering::Type ordering = m_fem.displacementFes->GetOrdering();

        for (const ElementPAData &data : m_elements) {
            MFEM_VERIFY(data.integrationRule != nullptr, "Prepared gravity force has no integration rule.");

            m_densityVariationLocal.GetSubVector(data.densityDofs, m_elementDensityVariation);
            m_gravityGradientVariationLocal.GetSubVector(data.gravityGradientDofs, m_elementGravityGradientVariation);
            m_displacementVariationLocal.GetSubVector(data.displacementDofs, m_elementDisplacementVariation);
            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(m_elementDensityVariation);
            }
            if (data.gravityGradientDofTransformation != nullptr) {
                data.gravityGradientDofTransformation->InvTransformPrimal(m_elementGravityGradientVariation);
            }
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(m_elementDisplacementVariation);
            }

            const mfem::FiniteElement &densityElement         = *m_fem.densityFes->GetFE(data.elementId);
            const mfem::FiniteElement &gravityGradientElement = *m_fem.gravityFluxFes->GetFE(data.elementId);
            const mfem::FiniteElement &displacementElement    = *m_fem.displacementFes->GetFE(data.elementId);
            mfem::ElementTransformation *transformation       = m_fem.mesh->GetElementTransformation(data.elementId);
            MFEM_VERIFY(transformation != nullptr, "Prepared gravity force received a null transformation.");

            const mapping::ElementDisplacementData directionData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, m_elementDisplacementVariation);
            const mfem::DenseMatrix &directionDofs = directionData.GetDofMatrix();
            const int scalarDisplacementDofCount   = displacementElement.GetDof();

            m_densityShape.SetSize(densityElement.GetDof());
            m_displacementShape.SetSize(scalarDisplacementDofCount);
            m_gravityGradientShape.SetSize(gravityGradientElement.GetDof(), dimension);
            m_referenceDisplacementDShape.SetSize(scalarDisplacementDofCount, dimension);
            m_referenceDisplacementJacobian.SetSize(dimension, dimension);
            m_displacementJacobianVariation.SetSize(dimension, dimension);
            m_mappingJacobian.SetSize(dimension, dimension);
            m_inverseMeshJacobian.SetSize(dimension, dimension);
            m_baseGravityReferenceValue.SetSize(dimension);
            m_gravityVariationReferenceValue.SetSize(dimension);
            m_mappedBaseGravity.SetSize(dimension);
            m_mappedGravityVariation.SetSize(dimension);
            m_mappedGeometryVariation.SetSize(dimension);
            m_forceValue.SetSize(dimension);
            m_elementAction.SetSize(data.displacementDofs.Size());
            m_elementAction = 0.0;

            for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
                densityElement.CalcShape(integrationPoint, m_densityShape);
                displacementElement.CalcShape(integrationPoint, m_displacementShape);
                displacementElement.CalcDShape(integrationPoint, m_referenceDisplacementDShape);
                mfem::MultAtB(directionDofs, m_referenceDisplacementDShape, m_referenceDisplacementJacobian);

                transformation->SetIntPoint(&integrationPoint);
                gravityGradientElement.CalcVShape(*transformation, m_gravityGradientShape);
                m_gravityGradientShape.MultTranspose(
                    m_elementGravityGradientVariation, m_gravityVariationReferenceValue
                );

                for (int row = 0; row < dimension; ++row) {
                    m_baseGravityReferenceValue(row) = data.baseGravityReferenceValues(quadraturePoint, row);
                    for (int column = 0; column < dimension; ++column) {
                        const int entry                    = row * dimension + column;
                        m_mappingJacobian(row, column)     = data.mappingJacobians(quadraturePoint, entry);
                        m_inverseMeshJacobian(row, column) = data.inverseMeshJacobians(quadraturePoint, entry);
                    }
                }
                mfem::Mult(m_referenceDisplacementJacobian, m_inverseMeshJacobian, m_displacementJacobianVariation);
                m_mappingJacobian.Mult(m_baseGravityReferenceValue, m_mappedBaseGravity);
                m_mappingJacobian.Mult(m_gravityVariationReferenceValue, m_mappedGravityVariation);
                m_displacementJacobianVariation.Mult(m_baseGravityReferenceValue, m_mappedGeometryVariation);

                const double densityVariationValue = m_elementDensityVariation * m_densityShape;
                const double baseDensityValue      = data.baseDensityValues(quadraturePoint);
                m_forceValue                       = 0.0;
                m_forceValue.Add(densityVariationValue, m_mappedBaseGravity);
                m_forceValue.Add(baseDensityValue, m_mappedGravityVariation);
                m_forceValue.Add(baseDensityValue, m_mappedGeometryVariation);
                m_forceValue *= data.referenceWeights(quadraturePoint);

                for (int scalarDof = 0; scalarDof < scalarDisplacementDofCount; ++scalarDof) {
                    for (int component = 0; component < dimension; ++component) {
                        const int vectorDof =
                            vector_dof_index(ordering, scalarDof, component, scalarDisplacementDofCount, dimension);
                        m_elementAction(vectorDof) += m_displacementShape(scalarDof) * m_forceValue(component);
                    }
                }
            }

            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->TransformDual(m_elementAction);
            }
            m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
        }

        local_to_true(*m_fem.displacementFes, m_localAction, actionTrue);
    }

    void PreparedGravityDisplacementForceOperator::ApplyCompleteJacobianAction(
        const mfem::Vector &densityVariation,
        const mfem::Vector &displacementVariation,
        const mfem::Vector &gravityGradientVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
        m_gravityGradientVariationTrue.SetSize(m_gravityContext.GetGravityGradientMap().full_size());
        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_gravityContext.GetGravityGradientMap().scatter(gravityGradientVariation, m_gravityGradientVariationTrue);
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        ApplyPreparedCompleteJacobianActionTrue(
            m_densityVariationTrue, m_displacementVariationTrue, m_gravityGradientVariationTrue, m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_densityJacobianStatistics.applications;
        ++m_gravityGradientJacobianStatistics.applications;
        ++m_displacementJacobianStatistics.applications;
        ++m_completeJacobianStatistics.applications;
    }

    bool PreparedGravityDisplacementForceOperator::IsPrepared() const noexcept {
        if (!m_isPrepared || !m_gravityContext.IsPrepared()) {
            return false;
        }

        return relevant_revisions_match(m_gravityContext.GetRevisions(), m_preparedRevisions);
    }

    std::uint64_t PreparedGravityDisplacementForceOperator::GetResidualPreparationCount() const noexcept {
        return m_residualPreparationCount;
    }

    std::uint64_t PreparedGravityDisplacementForceOperator::GetResidualApplicationCount() const noexcept {
        return m_residualApplicationCount;
    }

    const PreparedGravityDisplacementForceColumnStatistics &
    PreparedGravityDisplacementForceOperator::GetDensityJacobianStatistics() const noexcept {
        return m_densityJacobianStatistics;
    }

    const PreparedGravityDisplacementForceColumnStatistics &
    PreparedGravityDisplacementForceOperator::GetGravityGradientJacobianStatistics() const noexcept {
        return m_gravityGradientJacobianStatistics;
    }

    const PreparedGravityDisplacementForceColumnStatistics &
    PreparedGravityDisplacementForceOperator::GetDisplacementJacobianStatistics() const noexcept {
        return m_displacementJacobianStatistics;
    }

    const PreparedGravityDisplacementForceCompleteStatistics &
    PreparedGravityDisplacementForceOperator::GetCompleteJacobianStatistics() const noexcept {
        return m_completeJacobianStatistics;
    }

    const fem::FEM &PreparedGravityDisplacementForceOperator::GetFEM() const noexcept {
        return m_fem;
    }

    const context::gravity_field::GravityFieldLinearizationContext &
    PreparedGravityDisplacementForceOperator::GetGravityContext() const noexcept {
        return m_gravityContext;
    }

    void PreparedGravityDisplacementForceOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedGravityDisplacementForceOperator must be prepared for "
                          "the current shared gravity-context revisions before residual or "
                          "Jacobian application."
        );
    }

    PreparedGravityDisplacementForceJacobianOperator::PreparedGravityDisplacementForceJacobianOperator(
        const GravityDisplacementForceLayout &layout,
        const PreparedGravityDisplacementForceOperator &preparedOperator
    )
        : mfem::Operator(
              layout.residual_offsets().Last(),
              layout.value_offsets().Last()
          ),
          m_layout(layout),
          m_preparedOperator(preparedOperator) {
        using Form                  = utils::blocks::barotropic_equilibrium_form;

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

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

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

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

        MFEM_VERIFY(
            m_layout.size(densityValue) == m_preparedOperator.GetGravityContext().GetDensityMap().reduced_size() &&
                m_layout.size(displacementValue) ==
                    m_preparedOperator.GetGravityContext().GetDisplacementMap().reduced_size() &&
                m_layout.size(gravityGradientValue) ==
                    m_preparedOperator.GetGravityContext().GetGravityGradientMap().reduced_size() &&
                m_layout.size(displacementResidual) ==
                    m_preparedOperator.GetGravityContext().GetDisplacementMap().reduced_size(),
            "Prepared gravity-displacement-force MFEM adapter received "
            "incompatible coupled block sizes."
        );
    }

    void PreparedGravityDisplacementForceJacobianOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_preparedOperator.IsPrepared(), "Prepared gravity-displacement-force MFEM adapter requires a "
                                             "prepared operator."
        );

        MFEM_VERIFY(
            direction.Size() == Width(), "Prepared gravity-displacement-force MFEM adapter received a "
                                         "direction with the wrong size."
        );

        using Form                  = utils::blocks::barotropic_equilibrium_form;

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

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

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

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

        const mfem::Vector densityVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(densityValue), m_layout.size(densityValue)
        );

        const mfem::Vector displacementVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(displacementValue),
            m_layout.size(displacementValue)
        );

        const mfem::Vector gravityGradientVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(gravityGradientValue),
            m_layout.size(gravityGradientValue)
        );

        mfem::Vector displacementAction;

        m_preparedOperator.ApplyCompleteJacobianAction(
            densityVariation, displacementVariation, gravityGradientVariation, displacementAction
        );

        MFEM_VERIFY(
            displacementAction.Size() == m_layout.size(displacementResidual),
            "Prepared gravity-displacement-force MFEM adapter produced a "
            "displacement action with the wrong size."
        );

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

        const int residualOffset = m_layout.offset(displacementResidual);

        for (int entry = 0; entry < displacementAction.Size(); ++entry) {
            action(residualOffset + entry) = displacementAction(entry);
        }
    }

    const GravityDisplacementForceLayout &PreparedGravityDisplacementForceJacobianOperator::GetLayout() const noexcept {
        return m_layout;
    }
} // namespace mean_field::operators