module;

#include <array>
#include <cmath>
#include <cstdint>
#include <limits>
#include <utility>

#include <mfem.hpp>

module mean_field;

import :operators.prepared_pressure_force;
import :field.registry;
import :utils.blocks;
import :utils.domain;

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

    using PressureDomain = mean_field::field::FieldDomainT<mean_field::field::Enthalpy>;

    void verify_required_spaces(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(f.mesh != nullptr, "PreparedPressureForceOperator requires a mesh.");

        MFEM_VERIFY(
            f.enthalpyFes != nullptr, "PreparedPressureForceOperator requires the enthalpy "
                                      "finite-element space."
        );

        MFEM_VERIFY(
            f.displacementFes != nullptr, "PreparedPressureForceOperator requires the displacement "
                                          "finite-element space."
        );

        MFEM_VERIFY(
            f.compactificationFes != nullptr, "PreparedPressureForceOperator requires the compactification "
                                              "finite-element space."
        );

        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "PreparedPressureForceOperator requires the compactification "
                                                     "coordinate."
        );

        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "PreparedPressureForceOperator requires the quadrature-rule "
                                            "factory."
        );
    }

    [[nodiscard]]
    bool element_is_in_pressure_support(const int attribute) {
        return DomainSchema::template attribute_belongs_to<PressureDomain>(attribute);
    }

    void true_to_local(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &trueVector,
        mfem::Vector &localVector
    ) {
        MFEM_VERIFY(
            trueVector.Size() == finiteElementSpace.GetTrueVSize(),
            "Prepared pressure-force 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(),
            "Prepared pressure-force 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;
        }
    }

    [[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;
        }

        if (ordering == mfem::Ordering::byVDIM) {
            return scalarDof * dimension + component;
        }

        MFEM_ABORT(
            "The prepared pressure-force displacement space uses an "
            "unsupported ordering."
        );

        return -1;
    }

    [[nodiscard]]
    int get_pressure_extra_order(const mean_field::eos::Polytrope &equationOfState) {
        const double extraOrder =
            equationOfState.polytropic_index() * static_cast<double>(mean_field::field::Enthalpy::Scalar::familyOrder);

        MFEM_VERIFY(
            std::isfinite(extraOrder) && extraOrder >= 0.0 &&
                extraOrder <= static_cast<double>(std::numeric_limits<int>::max()),
            "The prepared pressure-force EOS effective polynomial order "
            "is invalid."
        );

        return static_cast<int>(std::ceil(extraOrder));
    }

    [[nodiscard]]
    const mfem::IntegrationRule &get_pressure_force_rule(
        const mean_field::fem::FEM &f,
        const mean_field::eos::Polytrope &equationOfState,
        const mfem::FiniteElement &enthalpyElement,
        const mfem::FiniteElement &displacementElement,
        const mfem::ElementTransformation &transformation
    ) {
        using EnthalpyField = mean_field::field::Field<mean_field::field::Enthalpy>;

        MFEM_VERIFY(
            enthalpyElement.GetOrder() == mean_field::field::Enthalpy::Scalar::familyOrder,
            "The prepared pressure-force enthalpy element does not "
            "match the registered enthalpy field."
        );

        MFEM_VERIFY(
            displacementElement.GetOrder() == mean_field::field::Displacement::Vector::familyOrder,
            "The prepared pressure-force test element does not match "
            "the registered displacement field."
        );

        /*
         * Query.domain remains legacy quadrature metadata for now.
         *
         * Physical element selection is no longer based on utils::DOMAINS;
         * it is performed from Enthalpy::Support + DomainSchema in
         * PrepareStaticPlan().
         */
        const mean_field::quadrature::Query query =
            EnthalpyField::make_query<mean_field::field::Enthalpy::Form::PressureForce>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(),
                std::array<int, 1>{get_pressure_extra_order(equationOfState)}, 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 prepared "
                                              "pressure-force integration rule."
        );

        return *rule.integration_rule;
    }
} // namespace

namespace mean_field::operators {
    struct PreparedPressureForceOperator::ConstructionData final {
        field::FieldDofMap enthalpyMap;

        field::FieldDofMap displacementMap;

        explicit ConstructionData(const fem::FEM &f)
            : enthalpyMap(
                  field::make_field_dof_map<
                      field::Enthalpy,
                      DomainSchema>(*f.enthalpyFes)
              ),
              displacementMap(
                  field::make_field_dof_map<
                      field::Displacement,
                      DomainSchema>(*f.displacementFes)
              ) {
        }
    };

    PreparedPressureForceOperator::ConstructionData
    PreparedPressureForceOperator::MakeConstructionData(const fem::FEM &f) {
        verify_required_spaces(f);

        return ConstructionData(f);
    }

    PreparedPressureForceOperator::PreparedPressureForceOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &equationOfState
    )
        : PreparedPressureForceOperator(
              f,
              domainMapper,
              equationOfState,
              MakeConstructionData(f)
          ) {
    }

    PreparedPressureForceOperator::PreparedPressureForceOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &equationOfState,
        ConstructionData constructionData
    )
        : m_fem(f),
          m_domainMapper(domainMapper),
          m_equationOfState(equationOfState),
          m_enthalpyMap(std::move(constructionData.enthalpyMap)),
          m_displacementMap(std::move(constructionData.displacementMap)),
          m_context(
              f,
              domainMapper,
              m_enthalpyMap,
              m_displacementMap
          ) {
        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_fem.mesh->Dimension(),
            "The prepared pressure-force mapper dimension does not "
            "match the mesh dimension."
        );

        MFEM_VERIFY(
            m_fem.displacementFes->GetVDim() == m_fem.mesh->Dimension(),
            "The prepared pressure-force displacement dimension does "
            "not match the mesh dimension."
        );

        MFEM_VERIFY(
            m_fem.displacementFes->GetOrdering() == mfem::Ordering::byNODES,
            "PreparedPressureForceOperator requires the registered "
            "byNODES displacement ordering."
        );

        MFEM_VERIFY(
            m_enthalpyMap.full_size() == m_fem.enthalpyFes->GetTrueVSize(),
            "The pressure-force enthalpy FieldDofMap does not match "
            "the enthalpy finite-element space."
        );

        MFEM_VERIFY(
            m_displacementMap.full_size() == m_fem.displacementFes->GetTrueVSize(),
            "The pressure-force displacement FieldDofMap does not "
            "match the displacement finite-element space."
        );

        m_baseEnthalpyTrue.SetSize(m_enthalpyMap.full_size());

        m_baseDisplacementTrue.SetSize(m_displacementMap.full_size());

        m_enthalpyVariationTrue.SetSize(m_enthalpyMap.full_size());

        m_displacementVariationTrue.SetSize(m_displacementMap.full_size());

        m_fullDisplacementAction.SetSize(m_displacementMap.full_size());

        m_baseEnthalpyTrue          = 0.0;

        m_baseDisplacementTrue      = 0.0;

        m_enthalpyVariationTrue     = 0.0;

        m_displacementVariationTrue = 0.0;

        m_fullDisplacementAction    = 0.0;
    }

    PreparedPressureForceReport PreparedPressureForceOperator::Prepare(
        const context::pressure_force::PressureForceStateView &state,
        const context::pressure_force::PressureForceDependencies &dependencies
    ) {
        PreparedPressureForceReport report;

        report.contextReport = m_context.Prepare(state, dependencies);

        if (!report.contextReport.DidAnyWork() && m_isPrepared) {
            return report;
        }

        /*
         * Canonical FieldDof -> MFEM expansion.
         *
         * Unsupported enthalpy true DOFs are set exactly to zero.
         * Displacement currently has an identity map but is intentionally
         * routed through the same abstraction.
         */
        m_enthalpyMap.scatter(m_context.GetBaseEnthalpy(), m_baseEnthalpyTrue);

        m_displacementMap.scatter(m_context.GetDisplacement(), m_baseDisplacementTrue);

        m_isPrepared = false;

        if (report.contextReport.preparedStaticDependencies) {
            PrepareStaticPlan();
        }

        if (report.contextReport.preparedGeometryState) {
            PrepareGeometry();
        }

        if (report.contextReport.preparedMaterialState) {
            PrepareMaterialState();

            FinalizeDisplacementJacobianPreparation();

            AssembleCachedResidual();

            ++m_residualPreparationCount;

            report.preparedEnthalpyJacobianData     = true;

            report.preparedDisplacementJacobianData = true;

            report.preparedResidual                 = true;
        }

        MFEM_VERIFY(
            !m_elements.empty(), "PreparedPressureForceOperator found no elements in the "
                                 "pressure-force field support."
        );

        MFEM_VERIFY(
            m_cachedResidual.Size() == m_displacementMap.reduced_size(),
            "The prepared pressure-force residual has the wrong "
            "supported displacement size."
        );

        m_isPrepared = true;

        return report;
    }

    void PreparedPressureForceOperator::PrepareStaticPlan() {
        m_elements.clear();

        m_elements.reserve(m_fem.mesh->GetNE());

        mfem::Vector enthalpyShape;
        mfem::DenseMatrix displacementDShape;

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

            MFEM_VERIFY(
                transformation != nullptr, "Prepared pressure-force static planning received a null "
                                           "element transformation."
            );

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

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

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

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

            MFEM_VERIFY(
                enthalpyElement.GetGeomType() == displacementElement.GetGeomType() &&
                    enthalpyElement.GetGeomType() == compactificationElement.GetGeomType() &&
                    enthalpyElement.GetGeomType() == transformation->GetGeometryType(),
                "Prepared pressure-force element geometries do not agree."
            );

            m_elements.emplace_back();

            ElementPAData &data            = m_elements.back();

            data.elementId                 = elementId;

            data.enthalpyDofTransformation = m_fem.enthalpyFes->GetElementDofs(elementId, data.enthalpyDofs);

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

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

            data.integrationRule = &get_pressure_force_rule(
                m_fem, m_equationOfState, enthalpyElement, displacementElement, *transformation
            );

            const int dimension                  = m_fem.mesh->Dimension();

            const int quadraturePointCount       = data.integrationRule->GetNPoints();

            const int enthalpyDofCount           = enthalpyElement.GetDof();

            const int scalarDisplacementDofCount = displacementElement.GetDof();

            MFEM_VERIFY(quadraturePointCount > 0, "The prepared pressure-force integration rule is empty.");

            MFEM_VERIFY(
                data.enthalpyDofs.Size() == enthalpyDofCount, "The prepared pressure-force enthalpy element has an "
                                                              "unexpected DOF count."
            );

            MFEM_VERIFY(
                data.displacementDofs.Size() == scalarDisplacementDofCount * dimension,
                "The prepared pressure-force displacement element has "
                "an unexpected vector DOF count."
            );

            data.enthalpyBasis.SetSize(quadraturePointCount, enthalpyDofCount);

            data.referenceTestGradients.resize(quadraturePointCount);

            data.physicalTestGradients.resize(quadraturePointCount);

            enthalpyShape.SetSize(enthalpyDofCount);

            displacementDShape.SetSize(scalarDisplacementDofCount, dimension);

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);

                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                displacementElement.CalcDShape(integrationPoint, displacementDShape);

                for (int enthalpyDof = 0; enthalpyDof < enthalpyDofCount; ++enthalpyDof) {
                    data.enthalpyBasis(quadraturePoint, enthalpyDof) = enthalpyShape(enthalpyDof);
                }

                data.referenceTestGradients[quadraturePoint] = displacementDShape;
            }
        }
    }

    void PreparedPressureForceOperator::PrepareGeometry() {
        mfem::Vector displacementLocal;

        true_to_local(*m_fem.displacementFes, m_baseDisplacementTrue, displacementLocal);

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

        mfem::Vector elementDisplacement;
        mfem::Vector elementCompactification;

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

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

            MFEM_VERIFY(
                transformation != nullptr, "Prepared pressure-force geometry received a null "
                                           "element transformation."
            );

            displacementLocal.GetSubVector(data.displacementDofs, elementDisplacement);

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

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

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

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

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

            data.baseDisplacementData.emplace(
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementDisplacement)
            );

            data.compactificationData.emplace(compactificationElement, elementCompactification);

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

            const int quadraturePointCount = data.integrationRule->GetNPoints();

            MFEM_VERIFY(
                static_cast<int>(data.referenceTestGradients.size()) == quadraturePointCount,
                "Prepared pressure-force geometry has inconsistent "
                "static gradient data."
            );

            data.quadratureWeights.SetSize(quadraturePointCount);

            data.baseMappingContexts.resize(quadraturePointCount);

            data.physicalTestGradients.resize(quadraturePointCount);

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);

                transformation->SetIntPoint(&integrationPoint);

                mapping::VolumeMappingContext &mappingContext = data.baseMappingContexts[quadraturePoint];

                const mapping::MappingStatus mappingStatus    = m_domainMapper.EvaluateVolume(
                    mappingData, *transformation, integrationPoint, workspace, mappingContext
                );

                MFEM_VERIFY(
                    mappingStatus == mapping::MappingStatus::valid,
                    "Stateless mapping failed while preparing "
                    "pressure-force geometry. Element: "
                        << data.elementId << ", attribute: " << transformation->Attribute
                        << ", quadrature point: " << quadraturePoint << ", status: " << static_cast<int>(mappingStatus)
                );

                const double quadratureWeight = mappingContext.quadrature.weight;

                MFEM_VERIFY(
                    std::isfinite(quadratureWeight) && quadratureWeight > 0.0,
                    "Prepared pressure-force geometry encountered an "
                    "invalid quadrature weight."
                );

                data.quadratureWeights(quadraturePoint)        = quadratureWeight;

                const mfem::DenseMatrix &referenceTestGradient = data.referenceTestGradients[quadraturePoint];

                mfem::DenseMatrix &physicalTestGradient        = data.physicalTestGradients[quadraturePoint];

                MFEM_VERIFY(
                    referenceTestGradient.Width() == mappingContext.quadrature.J_inv.Height() &&
                        mappingContext.quadrature.J_inv.Width() == m_fem.mesh->Dimension(),
                    "Prepared pressure-force geometry encountered "
                    "incompatible test-gradient and inverse-Jacobian "
                    "dimensions."
                );

                physicalTestGradient.SetSize(referenceTestGradient.Height(), mappingContext.quadrature.J_inv.Width());

                mfem::Mult(referenceTestGradient, mappingContext.quadrature.J_inv, physicalTestGradient);
            }
        }
    }

    void PreparedPressureForceOperator::PrepareMaterialState() {
        mfem::Vector enthalpyLocal;

        true_to_local(*m_fem.enthalpyFes, m_baseEnthalpyTrue, enthalpyLocal);

        mfem::Vector elementEnthalpy;
        mfem::Vector quadratureEnthalpy;

        const int dimension                             = m_fem.mesh->Dimension();

        const mfem::Ordering::Type displacementOrdering = m_fem.displacementFes->GetOrdering();

        for (ElementPAData &data : m_elements) {
            enthalpyLocal.GetSubVector(data.enthalpyDofs, elementEnthalpy);

            if (data.enthalpyDofTransformation != nullptr) {
                data.enthalpyDofTransformation->InvTransformPrimal(elementEnthalpy);
            }

            const int quadraturePointCount                 = data.enthalpyBasis.Height();

            const int enthalpyDofCount                     = data.enthalpyBasis.Width();

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

            const int scalarDisplacementDofCount           = displacementElement.GetDof();

            const int displacementDofCount                 = data.displacementDofs.Size();

            MFEM_VERIFY(
                data.quadratureWeights.Size() == quadraturePointCount &&
                    static_cast<int>(data.physicalTestGradients.size()) == quadraturePointCount &&
                    displacementDofCount == scalarDisplacementDofCount * dimension,
                "Prepared pressure-force material state has stale "
                "geometry data."
            );

            quadratureEnthalpy.SetSize(quadraturePointCount);

            data.enthalpyBasis.Mult(elementEnthalpy, quadratureEnthalpy);

            data.pressure.SetSize(quadraturePointCount);

            data.pressureDerivative.SetSize(quadraturePointCount);

            data.elementResidual.SetSize(displacementDofCount);

            data.elementResidual = 0.0;

            data.enthalpyJacobian.SetSize(displacementDofCount, enthalpyDofCount);

            data.enthalpyJacobian = 0.0;

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const double enthalpy = quadratureEnthalpy(quadraturePoint);

                const eos::SpecificEnthalpyValue specificEnthalpy{enthalpy};
                const double pressure =
                    eos::evaluate<eos::quantity::Pressure>(m_equationOfState, specificEnthalpy).value();

                const double pressureDerivative =
                    eos::partialDerivative<eos::quantity::Pressure, eos::quantity::SpecificEnthalpy>(
                        m_equationOfState, specificEnthalpy
                    )
                        .value();

                const double quadratureWeight = data.quadratureWeights(quadraturePoint);

                MFEM_VERIFY(
                    std::isfinite(pressure) && std::isfinite(pressureDerivative),
                    "Prepared pressure-force material state encountered "
                    "a non-finite EOS value."
                );

                data.pressure(quadraturePoint)                = pressure;

                data.pressureDerivative(quadraturePoint)      = pressureDerivative;

                const mfem::DenseMatrix &physicalTestGradient = data.physicalTestGradients[quadraturePoint];

                MFEM_VERIFY(
                    physicalTestGradient.Height() == scalarDisplacementDofCount &&
                        physicalTestGradient.Width() == dimension,
                    "Prepared pressure-force material state has an "
                    "invalid physical test-gradient matrix."
                );

                for (int scalarDof = 0; scalarDof < scalarDisplacementDofCount; ++scalarDof) {
                    for (int component = 0; component < dimension; ++component) {
                        const int vectorDof = vector_dof_index(
                            displacementOrdering, scalarDof, component, scalarDisplacementDofCount, dimension
                        );

                        const double weightedTestGradient =
                            quadratureWeight * physicalTestGradient(scalarDof, component);

                        data.elementResidual(vectorDof) -= pressure * weightedTestGradient;

                        for (int enthalpyDof = 0; enthalpyDof < enthalpyDofCount; ++enthalpyDof) {
                            data.enthalpyJacobian(vectorDof, enthalpyDof) -=
                                pressureDerivative * weightedTestGradient *
                                data.enthalpyBasis(quadraturePoint, enthalpyDof);
                        }
                    }
                }
            }
        }

        ++m_enthalpyJacobianStatistics.preparations;
    }

    void PreparedPressureForceOperator::FinalizeDisplacementJacobianPreparation() {
        const int dimension = m_fem.mesh->Dimension();

        for (const ElementPAData &data : m_elements) {
            const int quadraturePointCount = data.integrationRule->GetNPoints();

            MFEM_VERIFY(
                data.baseDisplacementData.has_value() && data.compactificationData.has_value() &&
                    static_cast<int>(data.baseMappingContexts.size()) == quadraturePointCount &&
                    static_cast<int>(data.referenceTestGradients.size()) == quadraturePointCount &&
                    static_cast<int>(data.physicalTestGradients.size()) == quadraturePointCount &&
                    data.quadratureWeights.Size() == quadraturePointCount &&
                    data.pressure.Size() == quadraturePointCount,
                "Prepared pressure-force displacement Jacobian has "
                "inconsistent frozen data."
            );

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                MFEM_VERIFY(
                    data.referenceTestGradients[quadraturePoint].Width() == dimension &&
                        data.physicalTestGradients[quadraturePoint].Width() == dimension,
                    "Prepared pressure-force displacement Jacobian has "
                    "a gradient with the wrong dimension."
                );
            }
        }

        ++m_displacementJacobianStatistics.preparations;
    }

    void PreparedPressureForceOperator::AssembleCachedResidual() {
        mfem::Vector localResidual(m_fem.displacementFes->GetVSize());

        localResidual = 0.0;

        mfem::Vector elementResidual;

        for (const ElementPAData &data : m_elements) {
            elementResidual = data.elementResidual;

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

            localResidual.AddElementVector(data.displacementDofs, elementResidual);
        }

        local_to_true(*m_fem.displacementFes, localResidual, m_fullDisplacementAction);

        m_cachedResidual.SetSize(m_displacementMap.reduced_size());

        /*
         * FieldDofMap::gather does not resize its destination.
         */
        m_displacementMap.gather(m_fullDisplacementAction, m_cachedResidual);
    }

    void PreparedPressureForceOperator::BuildResidual(mfem::Vector &residual) const {
        VerifyPrepared();

        residual = m_cachedResidual;

        ++m_residualApplicationCount;
    }

    void PreparedPressureForceOperator::ApplyEnthalpyJacobianAction(
        const mfem::Vector &enthalpyVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        MFEM_VERIFY(
            enthalpyVariation.Size() == m_enthalpyMap.reduced_size(),
            "Prepared pressure-force enthalpy variation has the wrong "
            "supported size."
        );

        m_enthalpyMap.scatter(enthalpyVariation, m_enthalpyVariationTrue);

        mfem::Vector enthalpyVariationLocal;

        true_to_local(*m_fem.enthalpyFes, m_enthalpyVariationTrue, enthalpyVariationLocal);

        mfem::Vector localAction(m_fem.displacementFes->GetVSize());

        localAction = 0.0;

        mfem::Vector elementVariation;
        mfem::Vector elementAction;

        for (const ElementPAData &data : m_elements) {
            enthalpyVariationLocal.GetSubVector(data.enthalpyDofs, elementVariation);

            if (data.enthalpyDofTransformation != nullptr) {
                data.enthalpyDofTransformation->InvTransformPrimal(elementVariation);
            }

            elementAction.SetSize(data.enthalpyJacobian.Height());

            data.enthalpyJacobian.Mult(elementVariation, elementAction);

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

            localAction.AddElementVector(data.displacementDofs, elementAction);
        }

        local_to_true(*m_fem.displacementFes, localAction, m_fullDisplacementAction);

        action.SetSize(m_displacementMap.reduced_size());

        m_displacementMap.gather(m_fullDisplacementAction, action);

        ++m_enthalpyJacobianStatistics.applications;
    }

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

        MFEM_VERIFY(
            displacementVariation.Size() == m_displacementMap.reduced_size(),
            "Prepared pressure-force displacement variation has the "
            "wrong supported size."
        );

        m_displacementMap.scatter(displacementVariation, m_displacementVariationTrue);

        mfem::Vector displacementVariationLocal;

        true_to_local(*m_fem.displacementFes, m_displacementVariationTrue, displacementVariationLocal);

        mfem::Vector localAction(m_fem.displacementFes->GetVSize());

        localAction = 0.0;

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

        mfem::Vector elementDisplacementVariation;
        mfem::Vector elementAction;

        mfem::DenseMatrix physicalTestGradientVariation;

        const int dimension                             = m_fem.mesh->Dimension();

        const mfem::Ordering::Type displacementOrdering = m_fem.displacementFes->GetOrdering();

        for (const ElementPAData &data : m_elements) {
            MFEM_VERIFY(
                data.baseDisplacementData.has_value() && data.compactificationData.has_value() &&
                    data.integrationRule != nullptr,
                "Prepared pressure-force displacement Jacobian has "
                "invalid frozen element data."
            );

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

            MFEM_VERIFY(
                transformation != nullptr, "Prepared pressure-force displacement Jacobian received "
                                           "a null element transformation."
            );

            displacementVariationLocal.GetSubVector(data.displacementDofs, elementDisplacementVariation);

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

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

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

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

            const int quadraturePointCount       = data.integrationRule->GetNPoints();

            const int scalarDisplacementDofCount = displacementElement.GetDof();

            MFEM_VERIFY(
                static_cast<int>(data.baseMappingContexts.size()) == quadraturePointCount &&
                    data.pressure.Size() == quadraturePointCount,
                "Prepared pressure-force displacement Jacobian has "
                "stale quadrature data."
            );

            elementAction.SetSize(data.displacementDofs.Size());

            elementAction = 0.0;

            physicalTestGradientVariation.SetSize(scalarDisplacementDofCount, dimension);

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);

                transformation->SetIntPoint(&integrationPoint);

                mapping::VolumeMappingVariation variation;

                const mapping::MappingStatus mappingStatus = m_domainMapper.EvaluateVolumeVariation(
                    mappingData, directionData, *transformation, integrationPoint,
                    data.baseMappingContexts[quadraturePoint], workspace, variation
                );

                MFEM_VERIFY(
                    mappingStatus == mapping::MappingStatus::valid,
                    "Stateless mapping variation failed while applying "
                    "the prepared pressure-force displacement Jacobian. "
                    "Element: "
                        << data.elementId << ", attribute: " << transformation->Attribute
                        << ", quadrature point: " << quadraturePoint << ", status: " << static_cast<int>(mappingStatus)
                );

                mfem::Mult(
                    data.referenceTestGradients[quadraturePoint], variation.inverse_element_jacobian_variation,
                    physicalTestGradientVariation
                );

                const mfem::DenseMatrix &physicalTestGradient = data.physicalTestGradients[quadraturePoint];

                for (int scalarDof = 0; scalarDof < scalarDisplacementDofCount; ++scalarDof) {
                    for (int component = 0; component < dimension; ++component) {
                        const int vectorDof = vector_dof_index(
                            displacementOrdering, scalarDof, component, scalarDisplacementDofCount, dimension
                        );

                        const double gradientWeightVariation =
                            data.quadratureWeights(quadraturePoint) *
                                physicalTestGradientVariation(scalarDof, component) +
                            variation.weight_variation * physicalTestGradient(scalarDof, component);

                        const double contribution = data.pressure(quadraturePoint) * gradientWeightVariation;

                        MFEM_VERIFY(
                            std::isfinite(gradientWeightVariation) && std::isfinite(contribution),
                            "Prepared pressure-force displacement "
                            "Jacobian encountered a non-finite "
                            "contribution."
                        );

                        elementAction(vectorDof) -= contribution;
                    }
                }
            }

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

            localAction.AddElementVector(data.displacementDofs, elementAction);
        }

        local_to_true(*m_fem.displacementFes, localAction, m_fullDisplacementAction);

        action.SetSize(m_displacementMap.reduced_size());

        m_displacementMap.gather(m_fullDisplacementAction, action);

        ++m_displacementJacobianStatistics.applications;
    }

    void PreparedPressureForceOperator::ApplyCompleteJacobianAction(
        const mfem::Vector &enthalpyVariation,
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        mfem::Vector displacementAction;

        ApplyEnthalpyJacobianAction(enthalpyVariation, action);

        ApplyDisplacementJacobianAction(displacementVariation, displacementAction);

        MFEM_VERIFY(
            action.Size() == displacementAction.Size(), "Prepared pressure-force complete Jacobian produced "
                                                        "incompatible column actions."
        );

        action += displacementAction;

        ++m_completeJacobianStatistics.applications;
    }

    bool PreparedPressureForceOperator::IsPrepared() const noexcept {
        return m_isPrepared && m_context.IsPrepared();
    }

    int PreparedPressureForceOperator::GetEnthalpySize() const noexcept {
        return m_enthalpyMap.reduced_size();
    }

    int PreparedPressureForceOperator::GetDisplacementSize() const noexcept {
        return m_displacementMap.reduced_size();
    }

    const context::pressure_force::PressureForceLinearizationContext &
    PreparedPressureForceOperator::GetContext() const noexcept {
        return m_context;
    }

    const context::pressure_force::PressureForcePreparationStatistics &
    PreparedPressureForceOperator::GetContextPreparationStatistics() const noexcept {
        return m_context.GetPreparationStatistics();
    }

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

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

    const PreparedPressureForceEnthalpyJacobianStatistics &
    PreparedPressureForceOperator::GetEnthalpyJacobianStatistics() const noexcept {
        return m_enthalpyJacobianStatistics;
    }

    const PreparedPressureForceDisplacementJacobianStatistics &
    PreparedPressureForceOperator::GetDisplacementJacobianStatistics() const noexcept {
        return m_displacementJacobianStatistics;
    }

    const PreparedPressureForceCompleteJacobianStatistics &
    PreparedPressureForceOperator::GetCompleteJacobianStatistics() const noexcept {
        return m_completeJacobianStatistics;
    }

    std::size_t PreparedPressureForceOperator::GetStellarElementCount() const noexcept {
        return m_elements.size();
    }

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

    void PreparedPressureForceOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedPressureForceOperator must be prepared before "
                          "residual or Jacobian application."
        );
    }

    PreparedPressureForceJacobianOperator::PreparedPressureForceJacobianOperator(
        const BarotropicEquilibriumLayout &layout,
        const PreparedPressureForceOperator &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 displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);

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

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

        /*
         * This adapter consumes only d and h and contributes only R_d.
         *
         * Do not impose GetTrueVSize() assumptions on unrelated root
         * blocks. In particular rho and h may now be reduced FieldDof
         * coordinates.
         */
        MFEM_VERIFY(
            m_layout.size(displacementValue) == m_preparedOperator.GetDisplacementSize(),
            "Prepared pressure-force MFEM adapter received an "
            "incompatible displacement value block."
        );

        MFEM_VERIFY(
            m_layout.size(enthalpyValue) == m_preparedOperator.GetEnthalpySize(),
            "Prepared pressure-force MFEM adapter received an "
            "incompatible enthalpy value block."
        );

        MFEM_VERIFY(
            m_layout.size(displacementResidual) == m_preparedOperator.GetDisplacementSize(),
            "Prepared pressure-force MFEM adapter received an "
            "incompatible displacement residual block."
        );

        MFEM_VERIFY(
            Height() == m_layout.residual_offsets().Last() && Width() == m_layout.value_offsets().Last(),
            "Prepared pressure-force MFEM adapter has inconsistent "
            "operator dimensions."
        );
    }

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

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

        using Form = utils::blocks::barotropic_equilibrium_form;

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

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

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

        /*
         * MFEM does not provide a const non-owning Vector view.
         * These alias the packed direction but are passed only through
         * const references.
         */
        const mfem::Vector displacementVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(displacementValue),
            m_layout.size(displacementValue)
        );

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

        mfem::Vector displacementAction;

        m_preparedOperator.ApplyCompleteJacobianAction(enthalpyVariation, displacementVariation, displacementAction);

        MFEM_VERIFY(
            displacementAction.Size() == m_layout.size(displacementResidual),
            "Prepared pressure-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 BarotropicEquilibriumLayout &PreparedPressureForceJacobianOperator::GetLayout() const noexcept {
        return m_layout;
    }
} // namespace mean_field::operators