module;

#include <array>
#include <cmath>
#include <mfem.hpp>

module mean_field;

import :operators.prepared_mass_normalization;

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

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

    void validate_finite_vector(
        const mfem::Vector &vector,
        const char *message
    ) {
        for (int index = 0; index < vector.Size(); ++index) {
            MFEM_VERIFY(std::isfinite(vector(index)), message);
        }
    }

    void true_to_local(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &trueVector,
        mfem::Vector &localVector
    ) {
        MFEM_VERIFY(trueVector.Size() == finiteElementSpace.GetTrueVSize(), "True vector has the wrong size.");

        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
    ) {
        MFEM_VERIFY(localVector.Size() == finiteElementSpace.GetVSize(), "Local vector has the wrong size.");

        trueVector.SetSize(finiteElementSpace.GetTrueVSize());
        trueVector                         = 0.0;

        const mfem::Operator *prolongation = finiteElementSpace.GetProlongationMatrix();

        if (prolongation != nullptr) {
            prolongation->MultTranspose(localVector, trueVector);
        } else {
            trueVector = localVector;
        }
    }

    const mfem::IntegrationRule &get_mass_normalization_rule(
        const mean_field::fem::FEM &f,
        const mfem::FiniteElement &densityElement,
        const mfem::ElementTransformation &transformation
    ) {
        using DensityField = mean_field::field::Field<mean_field::field::Density>;

        MFEM_VERIFY(
            densityElement.GetOrder() == mean_field::field::Density::Scalar::familyOrder,
            "The mass-normalization element does not match the registered "
            "density field."
        );

        const mean_field::quadrature::Query query =
            DensityField::make_query<mean_field::field::Density::Form::MassNormalization>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), std::array<int, 0>{},
                mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
            );

        const auto resolution = f.quadratureFactory->get(query, transformation.GetGeometryType());

        MFEM_VERIFY(
            resolution.integration_rule != nullptr, "The quadrature policy did not return a mass-normalization rule."
        );

        return *resolution.integration_rule;
    }

    void validate_shared_gravity_revisions(
        const mean_field::operators::context::gravity_field::GravityFieldLinearizationContext &gravityContext,
        const mean_field::operators::MassNormalizationDependencies &dependencies
    ) {
        MFEM_VERIFY(
            gravityContext.IsPrepared(), "PreparedMassNormalizationOperator requires the shared gravity "
                                         "linearization context to be prepared first."
        );

        const auto &revisions = gravityContext.GetRevisions();

        MFEM_VERIFY(
            revisions.discretization.value == dependencies.discretization.revision &&
                revisions.density.value == dependencies.density.revision &&
                revisions.displacement.value == dependencies.displacement.revision,
            "PreparedMassNormalizationOperator received dependency revisions "
            "that do not match the shared gravity context."
        );
    }

    void validate_shared_identity_transition(
        const mean_field::operators::MassNormalizationDependencyStamp &prepared,
        const mean_field::operators::MassNormalizationDependencyStamp &requested,
        const char *message
    ) {
        MFEM_VERIFY(prepared.identity == requested.identity || prepared.revision != requested.revision, message);
    }
} // namespace

namespace mean_field::operators {
    PreparedMassNormalizationOperator::PreparedMassNormalizationOperator(
        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, "PreparedMassNormalizationOperator requires a mesh.");
        MFEM_VERIFY(
            m_fem.densityFes != nullptr && m_fem.displacementFes != nullptr && m_fem.compactificationFes != nullptr &&
                m_fem.compactificationCoordinate != nullptr && m_fem.quadratureFactory != nullptr,
            "PreparedMassNormalizationOperator requires density, "
            "displacement, compactification, and quadrature data."
        );
        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_fem.mesh->Dimension(),
            "PreparedMassNormalizationOperator received a mapper with the "
            "wrong dimension."
        );

        MFEM_VERIFY(
            m_gravityContext.GetDensityMap().full_size() == m_fem.densityFes->GetTrueVSize() &&
                m_gravityContext.GetDisplacementMap().full_size() == m_fem.displacementFes->GetTrueVSize(),
            "PreparedMassNormalizationOperator received incompatible shared "
            "FieldDof maps."
        );

        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
    }

    PreparedMassNormalizationReport PreparedMassNormalizationOperator::Prepare(
        const MassNormalizationStateView &state,
        const MassNormalizationDependencies &dependencies
    ) {
        MFEM_VERIFY(
            std::isfinite(state.targetMass) && state.targetMass > 0.0,
            "PreparedMassNormalizationOperator requires a finite, positive "
            "target mass."
        );

        validate_shared_gravity_revisions(m_gravityContext, dependencies);

        if (m_isPrepared) {
            validate_shared_identity_transition(
                m_preparedDependencies.discretization, dependencies.discretization,
                "A new mass-normalization discretization identity must also "
                "change the shared gravity revision."
            );
            validate_shared_identity_transition(
                m_preparedDependencies.density, dependencies.density,
                "A new mass-normalization density identity must also change "
                "the shared gravity revision."
            );
            validate_shared_identity_transition(
                m_preparedDependencies.displacement, dependencies.displacement,
                "A new mass-normalization displacement identity must also "
                "change the shared gravity revision."
            );
        }

        const bool rebuildStaticPlan =
            !m_isPrepared || dependencies.discretization != m_preparedDependencies.discretization;

        const bool refreshGeometry =
            rebuildStaticPlan || dependencies.displacement != m_preparedDependencies.displacement;

        const bool refreshDensity   = rebuildStaticPlan || dependencies.density != m_preparedDependencies.density;

        const bool updateTargetMass = !m_isPrepared || dependencies.targetMass != m_preparedDependencies.targetMass ||
                                      state.targetMass != m_targetMass;

        m_isPrepared                = false;

        PreparedMassNormalizationReport report;

        if (rebuildStaticPlan) {
            BuildStaticPlan();
            report.rebuiltStaticPlan = true;
        }

        if (refreshGeometry) {
            RefreshGeometry(m_gravityContext.GetGeometryContext().GetDisplacementTrue());
            report.refreshedGeometry = true;
        }

        if (refreshDensity) {
            RefreshDensity(m_gravityContext.GetDensityTrue());
            report.refreshedDensity = true;
        }

        if (updateTargetMass) {
            m_targetMass             = state.targetMass;
            report.updatedTargetMass = true;
        }

        if (refreshGeometry || refreshDensity) {
            AssembleResidual();
            report.assembledResidual = true;
        } else if (updateTargetMass) {
            m_cachedResidual.SetSize(1);
            m_cachedResidual(0) = m_currentMass - m_targetMass;
            ++m_preparationCount;
            report.assembledResidual = true;
        }

        m_preparedDependencies = dependencies;
        m_isPrepared           = true;
        return report;
    }

    PreparedMassNormalizationReport PreparedMassNormalizationOperator::Prepare(
        const models::CompiledFixedMass &constraint,
        const MassNormalizationDependencies &dependencies
    ) {
        return Prepare({.targetMass = constraint.targetMass().value()}, dependencies);
    }

    void PreparedMassNormalizationOperator::BuildStaticPlan() {
        m_elements.clear();
        m_elements.reserve(m_fem.mesh->GetNE());

        int localStellarElementCount = 0;

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

            MFEM_VERIFY(
                transformation != nullptr, "PreparedMassNormalizationOperator received a null element "
                                           "transformation."
            );

            if (is_vacuum_attribute(transformation->Attribute)) {
                continue;
            }

            ++localStellarElementCount;
            m_elements.emplace_back();
            ElementPAData &data           = m_elements.back();
            data.elementId                = elementId;

            data.densityDofTransformation = m_fem.densityFes->GetElementDofs(elementId, data.densityDofs);

            data.displacementDofTransformation =
                m_fem.displacementFes->GetElementVDofs(elementId, data.displacementDofs);

            data.compactificationDofTransformation =
                m_fem.compactificationFes->GetElementDofs(elementId, data.compactificationDofs);

            const mfem::FiniteElement &densityElement = *m_fem.densityFes->GetFE(elementId);

            const mfem::IntegrationRule &integrationRule =
                get_mass_normalization_rule(m_fem, densityElement, *transformation);

            data.quadraturePoints.resize(integrationRule.GetNPoints());

            for (int quadraturePoint = 0; quadraturePoint < integrationRule.GetNPoints(); ++quadraturePoint) {
                QuadraturePointData &point = data.quadraturePoints[quadraturePoint];

                point.integrationPoint     = integrationRule.IntPoint(quadraturePoint);

                point.densityShape.SetSize(densityElement.GetDof());
                densityElement.CalcShape(point.integrationPoint, point.densityShape);
            }
        }

        int globalStellarElementCount = 0;
        MPI_Allreduce(
            &localStellarElementCount, &globalStellarElementCount, 1, MPI_INT, MPI_SUM, m_fem.mesh->GetComm()
        );

        MFEM_VERIFY(globalStellarElementCount > 0, "PreparedMassNormalizationOperator found no stellar elements.");
    }

    void PreparedMassNormalizationOperator::RefreshGeometry(const mfem::Vector &displacement) {
        MFEM_VERIFY(
            displacement.Size() == m_fem.displacementFes->GetTrueVSize(),
            "PreparedMassNormalizationOperator received a displacement "
            "vector with the wrong size."
        );
        validate_finite_vector(
            displacement, "PreparedMassNormalizationOperator received a non-finite "
                          "displacement value."
        );

        mfem::Vector displacementLocal;
        true_to_local(*m_fem.displacementFes, displacement, displacementLocal);

        mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());

        for (ElementPAData &data : m_elements) {
            displacementLocal.GetSubVector(data.displacementDofs, data.baseDisplacement);

            m_fem.compactificationCoordinate->GetSubVector(data.compactificationDofs, data.compactification);

            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(data.baseDisplacement);
            }

            if (data.compactificationDofTransformation != nullptr) {
                data.compactificationDofTransformation->InvTransformPrimal(data.compactification);
            }

            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(data.elementId);

            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);

            const mapping::ElementDisplacementData displacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);

            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, data.compactification
            );

            const mapping::ElementMappingData mappingData{
                .displacement = displacementData, .compactification = compactificationData
            };

            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);

            for (QuadraturePointData &point : data.quadraturePoints) {
                const mapping::MappingStatus status = m_domainMapper.EvaluateVolume(
                    mappingData, *transformation, point.integrationPoint, workspace, point.mappingContext
                );

                MFEM_VERIFY(
                    status == mapping::MappingStatus::valid, "Stateless mapping failed while preparing mass "
                                                             "normalization. Element: "
                                                                 << data.elementId
                                                                 << ", attribute: " << transformation->Attribute
                                                                 << ", status: " << static_cast<int>(status)
                );
            }
        }
    }

    void PreparedMassNormalizationOperator::RefreshDensity(const mfem::Vector &density) {
        MFEM_VERIFY(
            density.Size() == m_fem.densityFes->GetTrueVSize(),
            "PreparedMassNormalizationOperator received a density vector "
            "with the wrong size."
        );
        validate_finite_vector(
            density, "PreparedMassNormalizationOperator received a non-finite density "
                     "value."
        );

        mfem::Vector densityLocal;
        true_to_local(*m_fem.densityFes, density, densityLocal);

        mfem::Vector elementDensity;

        for (ElementPAData &data : m_elements) {
            densityLocal.GetSubVector(data.densityDofs, elementDensity);

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(elementDensity);
            }

            for (QuadraturePointData &point : data.quadraturePoints) {
                point.density = elementDensity * point.densityShape;
                MFEM_VERIFY(
                    std::isfinite(point.density), "PreparedMassNormalizationOperator produced a non-finite "
                                                  "quadrature density."
                );
            }
        }
    }

    void PreparedMassNormalizationOperator::AssembleResidual() {
        double localMass = 0.0;

        for (const ElementPAData &data : m_elements) {
            for (const QuadraturePointData &point : data.quadraturePoints) {
                localMass += point.density * point.mappingContext.quadrature.weight;
            }
        }

        m_currentMass = GlobalSum(localMass);
        MFEM_VERIFY(std::isfinite(m_currentMass), "PreparedMassNormalizationOperator assembled a non-finite mass.");

        m_cachedResidual.SetSize(1);
        m_cachedResidual(0) = m_currentMass - m_targetMass;
        ++m_preparationCount;
    }

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

    double PreparedMassNormalizationOperator::EvaluateDensityActionLocal(const mfem::Vector &densityVariation) const {
        MFEM_VERIFY(
            densityVariation.Size() == m_fem.densityFes->GetTrueVSize(),
            "Mass-normalization density action received a vector with the "
            "wrong size."
        );
        validate_finite_vector(densityVariation, "Mass-normalization density action received a non-finite value.");

        mfem::Vector densityVariationLocal;
        true_to_local(*m_fem.densityFes, densityVariation, densityVariationLocal);

        mfem::Vector elementDensityVariation;
        double localAction = 0.0;

        for (const ElementPAData &data : m_elements) {
            densityVariationLocal.GetSubVector(data.densityDofs, elementDensityVariation);

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(elementDensityVariation);
            }

            for (const QuadraturePointData &point : data.quadraturePoints) {
                localAction += (elementDensityVariation * point.densityShape) * point.mappingContext.quadrature.weight;
            }
        }

        return localAction;
    }

    double PreparedMassNormalizationOperator::EvaluateDisplacementActionLocal(
        const mfem::Vector &displacementVariation
    ) const {
        MFEM_VERIFY(
            displacementVariation.Size() == m_fem.displacementFes->GetTrueVSize(),
            "Mass-normalization displacement action received a vector with "
            "the wrong size."
        );
        validate_finite_vector(
            displacementVariation, "Mass-normalization displacement action received a non-finite "
                                   "value."
        );

        mfem::Vector displacementVariationLocal;
        true_to_local(*m_fem.displacementFes, displacementVariation, displacementVariationLocal);

        mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());
        mapping::VolumeMappingVariation variation;

        mfem::Vector elementDisplacementVariation;
        double localAction = 0.0;

        for (const ElementPAData &data : m_elements) {
            displacementVariationLocal.GetSubVector(data.displacementDofs, elementDisplacementVariation);

            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(elementDisplacementVariation);
            }

            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(data.elementId);

            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);

            const mapping::ElementDisplacementData baseDisplacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);

            const mapping::ElementDisplacementData directionData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementDisplacementVariation);

            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, data.compactification
            );

            const mapping::ElementMappingData mappingData{
                .displacement = baseDisplacementData, .compactification = compactificationData
            };

            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);

            for (const QuadraturePointData &point : data.quadraturePoints) {
                const mapping::MappingStatus status = m_domainMapper.EvaluateVolumeVariation(
                    mappingData, directionData, *transformation, point.integrationPoint, point.mappingContext,
                    workspace, variation
                );

                MFEM_VERIFY(
                    status == mapping::MappingStatus::valid, "Stateless mapping variation failed in the "
                                                             "mass-normalization displacement action. Element: "
                                                                 << data.elementId
                                                                 << ", status: " << static_cast<int>(status)
                );

                localAction += point.density * variation.weight_variation;
            }
        }

        return localAction;
    }

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

        MFEM_VERIFY(
            densityVariation.Size() == m_gravityContext.GetDensityMap().reduced_size(),
            "Mass-normalization density action received a supported vector "
            "with the wrong size."
        );
        validate_finite_vector(densityVariation, "Mass-normalization density action received a non-finite value.");
        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);

        action.SetSize(1);
        action(0) = GlobalSum(EvaluateDensityActionLocal(m_densityVariationTrue));
        ++m_actionStatistics.densityApplications;
    }

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

        MFEM_VERIFY(
            displacementVariation.Size() == m_gravityContext.GetDisplacementMap().reduced_size(),
            "Mass-normalization displacement action received a supported "
            "vector with the wrong size."
        );
        validate_finite_vector(
            displacementVariation, "Mass-normalization displacement action received a non-finite value."
        );
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        action.SetSize(1);
        action(0) = GlobalSum(EvaluateDisplacementActionLocal(m_displacementVariationTrue));
        ++m_actionStatistics.displacementApplications;
    }

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

        MFEM_VERIFY(
            densityVariation.Size() == m_gravityContext.GetDensityMap().reduced_size(),
            "Mass-normalization complete action received a supported density "
            "vector with the wrong size."
        );
        MFEM_VERIFY(
            displacementVariation.Size() == m_gravityContext.GetDisplacementMap().reduced_size(),
            "Mass-normalization complete action received a supported "
            "displacement vector with the wrong size."
        );
        validate_finite_vector(densityVariation, "Mass-normalization complete action received a non-finite density.");
        validate_finite_vector(
            displacementVariation, "Mass-normalization complete action received a non-finite displacement."
        );

        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        const double localAction = EvaluateDensityActionLocal(m_densityVariationTrue) +
                                   EvaluateDisplacementActionLocal(m_displacementVariationTrue);

        action.SetSize(1);
        action(0) = GlobalSum(localAction);
        ++m_actionStatistics.completeApplications;
    }

    void PreparedMassNormalizationOperator::ApplyJacobian(
        const FixedMassJacobianInput &input,
        mfem::Vector &action
    ) const {
        ApplyCompleteJacobianAction(input.densityVariation, input.displacementVariation, action);
    }

    void PreparedMassNormalizationOperator::AssembleDensityTransposeAction(
        const double residualDual,
        mfem::Vector &densityDual
    ) const {
        mfem::Vector localDual(m_fem.densityFes->GetVSize());
        localDual = 0.0;

        mfem::Vector elementDual;

        for (const ElementPAData &data : m_elements) {
            elementDual.SetSize(data.densityDofs.Size());
            elementDual = 0.0;

            for (const QuadraturePointData &point : data.quadraturePoints) {
                elementDual.Add(residualDual * point.mappingContext.quadrature.weight, point.densityShape);
            }

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->TransformDual(elementDual);
            }

            localDual.AddElementVector(data.densityDofs, elementDual);
        }

        mfem::Vector trueDual;
        local_to_true(*m_fem.densityFes, localDual, trueDual);

        densityDual.SetSize(m_gravityContext.GetDensityMap().reduced_size());
        m_gravityContext.GetDensityMap().gather(trueDual, densityDual);
    }

    void PreparedMassNormalizationOperator::AssembleDisplacementTransposeAction(
        const double residualDual,
        mfem::Vector &displacementDual
    ) const {
        mfem::Vector localDual(m_fem.displacementFes->GetVSize());
        localDual = 0.0;

        mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());
        mapping::VolumeMappingVariation variation;
        mfem::Vector elementDirection;
        mfem::Vector elementDual;

        for (const ElementPAData &data : m_elements) {
            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(data.elementId);
            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);

            const mapping::ElementDisplacementData baseDisplacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, data.compactification
            );
            const mapping::ElementMappingData mappingData{
                .displacement = baseDisplacementData, .compactification = compactificationData
            };

            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);

            elementDirection.SetSize(data.displacementDofs.Size());
            elementDual.SetSize(data.displacementDofs.Size());
            elementDual = 0.0;

            for (int elementDof = 0; elementDof < elementDirection.Size(); ++elementDof) {
                elementDirection             = 0.0;
                elementDirection(elementDof) = 1.0;

                const mapping::ElementDisplacementData directionData =
                    mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementDirection);

                double elementDofAction = 0.0;

                for (const QuadraturePointData &point : data.quadraturePoints) {
                    const mapping::MappingStatus status = m_domainMapper.EvaluateVolumeVariation(
                        mappingData, directionData, *transformation, point.integrationPoint, point.mappingContext,
                        workspace, variation
                    );

                    MFEM_VERIFY(
                        status == mapping::MappingStatus::valid,
                        "Stateless mapping variation failed in the mass-normalization transpose action. Element: "
                            << data.elementId << ", status: " << static_cast<int>(status)
                    );

                    elementDofAction += point.density * variation.weight_variation;
                }

                elementDual(elementDof) = residualDual * elementDofAction;
            }

            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->TransformDual(elementDual);
            }

            localDual.AddElementVector(data.displacementDofs, elementDual);
        }

        mfem::Vector trueDual;
        local_to_true(*m_fem.displacementFes, localDual, trueDual);

        displacementDual.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(trueDual, displacementDual);
    }

    void PreparedMassNormalizationOperator::ApplyCompleteJacobianTransposeAction(
        const double residualDual,
        mfem::Vector &densityDual,
        mfem::Vector &displacementDual
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(std::isfinite(residualDual), "Mass-normalization transpose action received a non-finite dual.");

        AssembleDensityTransposeAction(residualDual, densityDual);
        AssembleDisplacementTransposeAction(residualDual, displacementDual);
        ++m_actionStatistics.transposeApplications;
    }

    void PreparedMassNormalizationOperator::ApplyJacobianTranspose(
        const mfem::Vector &residualDual,
        FixedMassJacobianTransposeOutput output
    ) const {
        MFEM_VERIFY(residualDual.Size() == 1, "Fixed-mass transpose action requires one residual dual value.");
        ApplyCompleteJacobianTransposeAction(residualDual(0), output.densityDual, output.displacementDual);
    }

    double PreparedMassNormalizationOperator::GlobalSum(const double localValue) const {
        double globalValue = 0.0;
        MPI_Allreduce(&localValue, &globalValue, 1, MPI_DOUBLE, MPI_SUM, m_fem.mesh->GetComm());
        return globalValue;
    }

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

        const auto &revisions = m_gravityContext.GetRevisions();
        return revisions.discretization.value == m_preparedDependencies.discretization.revision &&
               revisions.density.value == m_preparedDependencies.density.revision &&
               revisions.displacement.value == m_preparedDependencies.displacement.revision;
    }

    double PreparedMassNormalizationOperator::GetCurrentMass() const {
        VerifyPrepared();
        return m_currentMass;
    }

    double PreparedMassNormalizationOperator::GetTargetMass() const {
        VerifyPrepared();
        return m_targetMass;
    }

    std::uint64_t PreparedMassNormalizationOperator::GetPreparationCount() const noexcept {
        return m_preparationCount;
    }

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

    const PreparedMassNormalizationActionStatistics &
    PreparedMassNormalizationOperator::GetActionStatistics() const noexcept {
        return m_actionStatistics;
    }

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

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

    void PreparedMassNormalizationOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedMassNormalizationOperator must be prepared for the "
                          "current shared gravity-context revisions."
        );
    }

    PreparedMassNormalizationJacobianOperator::PreparedMassNormalizationJacobianOperator(
        const MassNormalizationLayout &layout,
        const PreparedMassNormalizationOperator &preparedOperator
    )
        : mfem::Operator(
              layout.residual_offsets().Last(),
              layout.value_offsets().Last()
          ),
          m_layout(layout),
          m_preparedOperator(preparedOperator) {
        const fem::FEM &f = m_preparedOperator.GetFEM();

        MFEM_VERIFY(
            f.densityFes != nullptr && f.displacementFes != nullptr && f.gravityFluxFes != nullptr &&
                f.gravityPotentialFes != nullptr && f.enthalpyFes != nullptr,
            "Prepared mass-normalization MFEM adapter requires every "
            "finite-element space in the barotropic equilibrium layout."
        );

        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 gravityPotentialValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.poisson_term);
        constexpr auto enthalpyValue =
            utils::blocks::get_value_block<Form>(utils::blocks::enthalpy_field.specific_term);
        constexpr auto barotropicConstantValue =
            utils::blocks::get_value_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);
        constexpr auto gravityGradientResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.gradient_term);
        constexpr auto gravityPotentialResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.poisson_term);
        constexpr auto densityResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::density_field.mass_term);
        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);
        constexpr auto enthalpyResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term);
        constexpr auto massResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);

        using DomainSchema                   = utils::domain::CoreEnvelopeVacuumDomainSchema;

        const auto &gravityContext           = m_preparedOperator.GetGravityContext();

        const field::FieldDofMap enthalpyMap = field::make_field_dof_map<field::Enthalpy, DomainSchema>(*f.enthalpyFes);

        MFEM_VERIFY(
            m_layout.size(densityValue) == gravityContext.GetDensityMap().reduced_size() &&
                m_layout.size(displacementValue) == gravityContext.GetDisplacementMap().reduced_size() &&
                m_layout.size(gravityGradientValue) == gravityContext.GetGravityGradientMap().reduced_size() &&
                m_layout.size(gravityPotentialValue) == gravityContext.GetGravityPotentialMap().reduced_size() &&
                m_layout.size(enthalpyValue) == enthalpyMap.reduced_size() &&
                m_layout.size(barotropicConstantValue) == 1 &&
                m_layout.size(gravityGradientResidual) == gravityContext.GetGravityGradientMap().reduced_size() &&
                m_layout.size(gravityPotentialResidual) == gravityContext.GetGravityPotentialMap().reduced_size() &&
                m_layout.size(densityResidual) == gravityContext.GetDensityMap().reduced_size() &&
                m_layout.size(displacementResidual) == gravityContext.GetDisplacementMap().reduced_size() &&
                m_layout.size(enthalpyResidual) == enthalpyMap.reduced_size() && m_layout.size(massResidual) == 1,
            "Prepared mass-normalization MFEM adapter received incompatible "
            "barotropic block sizes."
        );
    }

    void PreparedMassNormalizationJacobianOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_preparedOperator.IsPrepared(), "Prepared mass-normalization MFEM adapter requires a prepared "
                                             "row operator."
        );
        MFEM_VERIFY(
            direction.Size() == Width(), "Prepared mass-normalization 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 massResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_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)
        );

        mfem::Vector massAction;
        m_preparedOperator.ApplyCompleteJacobianAction(densityVariation, displacementVariation, massAction);

        action.SetSize(Height());
        action                                = 0.0;
        action(m_layout.offset(massResidual)) = massAction(0);
    }

    void PreparedMassNormalizationJacobianOperator::MultTranspose(
        const mfem::Vector &residualDual,
        mfem::Vector &stateDual
    ) const {
        MFEM_VERIFY(
            m_preparedOperator.IsPrepared(),
            "Prepared mass-normalization MFEM adapter requires a prepared row operator."
        );
        MFEM_VERIFY(
            residualDual.Size() == Height(),
            "Prepared mass-normalization MFEM adapter received a residual dual 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 massResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);

        mfem::Vector densityDual;
        mfem::Vector displacementDual;
        m_preparedOperator.ApplyCompleteJacobianTransposeAction(
            residualDual(m_layout.offset(massResidual)), densityDual, displacementDual
        );

        stateDual.SetSize(Width());
        stateDual = 0.0;

        mfem::Vector densityBlock(stateDual.GetData() + m_layout.offset(densityValue), m_layout.size(densityValue));
        densityBlock = densityDual;

        mfem::Vector displacementBlock(
            stateDual.GetData() + m_layout.offset(displacementValue), m_layout.size(displacementValue)
        );
        displacementBlock = displacementDual;
    }

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