module;

#include <array>
#include <cmath>
#include <cstdint>
#include <expected>
#include <limits>
#include <optional>
#include <stdexcept>
#include <utility>

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

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

    using Rejection      = mean_field::operators::PressureForcePreparationRejection;
    using Reason         = mean_field::operators::PressureForcePreparationRejectionReason;

    [[nodiscard]] Rejection equation_of_state_rejection(const mean_field::eos::EvaluationErrorCode code) noexcept {
        return {.reason = Reason::equation_of_state, .equationOfStateCode = code};
    }

    [[nodiscard]] Rejection mapping_rejection(const mean_field::mapping::MappingStatus status) {
        MFEM_VERIFY(
            status != mean_field::mapping::MappingStatus::invalid_dimension,
            "Prepared pressure-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;
        }
        switch (rejection->reason) {
        case Reason::equation_of_state:
            return static_cast<int>(rejection->equationOfStateCode) + 1;
        case Reason::invalid_mapping:
            return 128 + static_cast<int>(rejection->mappingStatus);
        case Reason::non_finite_arithmetic:
        default:
            return 256;
        }
    }

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

    [[nodiscard]] std::optional<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("PreparedPressureForceOperator could not synchronize candidate validity.");
        }
        if (globalEncoded == 0) {
            return std::nullopt;
        }
        return decode_rejection(globalEncoded);
    }

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

    [[nodiscard]] bool matrix_is_finite(const mfem::DenseMatrix &matrix) noexcept {
        for (int row = 0; row < matrix.Height(); ++row) {
            for (int column = 0; column < matrix.Width(); ++column) {
                if (!std::isfinite(matrix(row, column))) {
                    return false;
                }
            }
        }
        return true;
    }

    [[noreturn]] void throw_rejection(const Rejection &rejection) {
        switch (rejection.reason) {
        case Reason::equation_of_state:
            throw mean_field::eos::EvaluationError(
                rejection.equationOfStateCode, "PreparedPressureForceOperator encountered invalid thermodynamic data."
            );
        case Reason::invalid_mapping:
            throw std::domain_error("PreparedPressureForceOperator encountered an invalid mapped domain.");
        case Reason::non_finite_arithmetic:
        default:
            throw std::domain_error("PreparedPressureForceOperator produced non-finite arithmetic.");
        }
    }

    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
    ) {
        auto result = TryPrepare(state, dependencies);
        if (!result.has_value()) {
            throw_rejection(result.error());
        }
        return std::move(result).value();
    }

    std::expected<
        PreparedPressureForceReport,
        PressureForcePreparationRejection>
    PreparedPressureForceOperator::TryPrepare(
        const context::pressure_force::PressureForceStateView &state,
        const context::pressure_force::PressureForceDependencies &dependencies
    ) {
        const bool wasPrepared = m_isPrepared;
        PreparedPressureForceReport report;

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

        if (!report.contextReport.DidAnyWork() && wasPrepared) {
            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 || !wasPrepared) {
            PrepareStaticPlan();
        }

        if (report.contextReport.preparedGeometryState || !wasPrepared) {
            const auto globalGeometryFailure = synchronize_rejection(PrepareGeometry(), m_fem.enthalpyFes->GetComm());
            if (globalGeometryFailure.has_value()) {
                return std::unexpected(*globalGeometryFailure);
            }
        }

        std::optional<PressureForcePreparationRejection> localMaterialFailure;
        if (report.contextReport.preparedMaterialState || !wasPrepared) {
            localMaterialFailure = PrepareMaterialState();
        }

        const auto globalMaterialFailure = synchronize_rejection(localMaterialFailure, m_fem.enthalpyFes->GetComm());
        if (globalMaterialFailure.has_value()) {
            return std::unexpected(*globalMaterialFailure);
        }

        if (report.contextReport.preparedMaterialState || !wasPrepared) {
            FinalizeDisplacementJacobianPreparation();

            const auto globalAssemblyFailure =
                synchronize_rejection(AssembleCachedResidual(), m_fem.enthalpyFes->GetComm());
            if (globalAssemblyFailure.has_value()) {
                return std::unexpected(*globalAssemblyFailure);
            }

            ++m_enthalpyJacobianStatistics.preparations;
            ++m_displacementJacobianStatistics.preparations;

            ++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());

        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."
            );

            const fem::ReferenceTableCache &referenceTables = m_fem.GetReferenceTables();
            data.enthalpyReferenceTable = referenceTables.GetScalarTable(enthalpyElement, *data.integrationRule);
            data.displacementReferenceTable =
                referenceTables.GetScalarTable(displacementElement, *data.integrationRule);
            data.physicalTestGradients.resize(quadraturePointCount);
        }
    }

    std::optional<PressureForcePreparationRejection> PreparedPressureForceOperator::PrepareGeometry() {
        mfem::Vector displacementLocal;

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

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

        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(
                data.displacementReferenceTable != nullptr &&
                    data.displacementReferenceTable->GetPointCount() == quadraturePointCount,
                "Prepared pressure-force geometry has inconsistent "
                "static gradient data."
            );

            data.quadratureWeights.SetSize(quadraturePointCount);

            data.baseMappingContexts.SetSize(quadraturePointCount, m_fem.mesh->Dimension());

            data.physicalTestGradients.resize(quadraturePointCount);

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

                transformation->SetIntPoint(&integrationPoint);

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

                if (mappingStatus != mapping::MappingStatus::valid) {
                    return mapping_rejection(mappingStatus);
                }

                const double quadratureWeight = mappingContext.quadrature.weight;

                if (!std::isfinite(quadratureWeight)) {
                    return mapping_rejection(mapping::MappingStatus::non_finite_result);
                }
                if (quadratureWeight <= 0.0) {
                    return mapping_rejection(mapping::MappingStatus::non_positive_determinant);
                }

                data.baseMappingContexts.Store(quadraturePoint, mappingContext);
                data.quadratureWeights(quadraturePoint) = quadratureWeight;

                const mfem::DenseMatrix &referenceTestGradient =
                    data.displacementReferenceTable->GetGradients(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);
                if (!matrix_is_finite(physicalTestGradient)) {
                    return non_finite_rejection();
                }
            }
        }
        return std::nullopt;
    }

    std::optional<PressureForcePreparationRejection> 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();
        std::optional<PressureForcePreparationRejection> materialFailure;

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

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

            const mfem::DenseMatrix &enthalpyBasis         = data.GetEnthalpyBasis();

            const int quadraturePointCount                 = enthalpyBasis.Height();

            const int enthalpyDofCount                     = 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);

            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);

                if (!std::isfinite(enthalpy)) {
                    materialFailure = equation_of_state_rejection(eos::EvaluationErrorCode::nonfinite_input);
                    data.pressure(quadraturePoint)           = 0.0;
                    data.pressureDerivative(quadraturePoint) = 0.0;
                    continue;
                }
                if (enthalpy < 0.0) {
                    materialFailure = equation_of_state_rejection(eos::EvaluationErrorCode::outside_domain);
                    data.pressure(quadraturePoint)           = 0.0;
                    data.pressureDerivative(quadraturePoint) = 0.0;
                    continue;
                }

                const dimensions::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);

                if (!std::isfinite(pressure) || !std::isfinite(pressureDerivative)) {
                    materialFailure = equation_of_state_rejection(eos::EvaluationErrorCode::nonfinite_result);
                    data.pressure(quadraturePoint)           = 0.0;
                    data.pressureDerivative(quadraturePoint) = 0.0;
                    continue;
                }

                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);

                        const double residualContribution = pressure * weightedTestGradient;
                        if (!std::isfinite(weightedTestGradient) || !std::isfinite(residualContribution)) {
                            materialFailure = non_finite_rejection();
                            continue;
                        }

                        data.elementResidual(vectorDof) -= residualContribution;

                        for (int enthalpyDof = 0; enthalpyDof < enthalpyDofCount; ++enthalpyDof) {
                            const double jacobianContribution =
                                pressureDerivative * weightedTestGradient * enthalpyBasis(quadraturePoint, enthalpyDof);
                            if (!std::isfinite(jacobianContribution)) {
                                materialFailure = non_finite_rejection();
                                continue;
                            }
                            data.enthalpyJacobian(vectorDof, enthalpyDof) -= jacobianContribution;
                        }
                    }
                }
            }

            if (!vector_is_finite(data.elementResidual) || !matrix_is_finite(data.enthalpyJacobian)) {
                materialFailure = non_finite_rejection();
            }
        }
        return materialFailure;
    }

    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() &&
                    data.baseMappingContexts.GetPointCount() == quadraturePointCount &&
                    data.baseMappingContexts.GetDimension() == dimension &&
                    data.displacementReferenceTable != nullptr &&
                    data.displacementReferenceTable->GetPointCount() == 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.displacementReferenceTable->GetGradients(quadraturePoint).Width() == dimension &&
                        data.physicalTestGradients[quadraturePoint].Width() == dimension,
                    "Prepared pressure-force displacement Jacobian has "
                    "a gradient with the wrong dimension."
                );
            }
        }
    }

    std::optional<PressureForcePreparationRejection> 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);
        if (!vector_is_finite(m_fullDisplacementAction) || !vector_is_finite(m_cachedResidual)) {
            return non_finite_rejection();
        }
        return std::nullopt;
    }

    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;

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

        mfem::DenseMatrix referenceDisplacementJacobian;
        mfem::DenseMatrix inverseElementJacobian;
        mfem::DenseMatrix inverseElementJacobianVariation;
        mfem::DenseMatrix matrixTemporary;
        mfem::DenseMatrix physicalTestGradientVariation;

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

        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 int quadraturePointCount                 = data.integrationRule->GetNPoints();

            const int scalarDisplacementDofCount           = displacementElement.GetDof();

            const mfem::DenseMatrix directionDofs(
                elementDisplacementVariation.GetData(), scalarDisplacementDofCount, dimension
            );

            MFEM_VERIFY(
                data.baseMappingContexts.GetPointCount() == quadraturePointCount &&
                    data.pressure.Size() == quadraturePointCount,
                "Prepared pressure-force displacement Jacobian has "
                "stale quadrature data."
            );

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

            elementAction = 0.0;

            referenceDisplacementJacobian.SetSize(dimension, dimension);
            inverseElementJacobianVariation.SetSize(dimension, dimension);
            matrixTemporary.SetSize(dimension, dimension);
            physicalTestGradientVariation.SetSize(scalarDisplacementDofCount, dimension);

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const mfem::DenseMatrix &referenceTestGradient =
                    data.displacementReferenceTable->GetGradients(quadraturePoint);

                mfem::MultAtB(directionDofs, referenceTestGradient, referenceDisplacementJacobian);

                data.baseMappingContexts.LoadInverseJacobian(quadraturePoint, inverseElementJacobian);
                mfem::Mult(inverseElementJacobian, referenceDisplacementJacobian, matrixTemporary);

                double logarithmicJacobianVariation{0.0};
                for (int component = 0; component < dimension; ++component) {
                    logarithmicJacobianVariation += matrixTemporary(component, component);
                }

                mfem::Mult(matrixTemporary, inverseElementJacobian, inverseElementJacobianVariation);
                inverseElementJacobianVariation *= -1.0;

                mfem::Mult(referenceTestGradient, inverseElementJacobianVariation, physicalTestGradientVariation);

                const mfem::DenseMatrix &physicalTestGradient = data.physicalTestGradients[quadraturePoint];
                const double quadratureWeight                 = data.quadratureWeights(quadraturePoint);
                const double pressure                         = data.pressure(quadraturePoint);

                for (int component = 0; component < dimension; ++component) {
                    const double *variationColumn =
                        physicalTestGradientVariation.GetData() + component * scalarDisplacementDofCount;
                    const double *physicalColumn =
                        physicalTestGradient.GetData() + component * scalarDisplacementDofCount;
                    double *actionColumn = elementAction.GetData() + component * scalarDisplacementDofCount;

                    for (int scalarDof = 0; scalarDof < scalarDisplacementDofCount; ++scalarDof) {
                        const double gradientWeightVariation =
                            quadratureWeight * variationColumn[scalarDof] +
                            quadratureWeight * logarithmicJacobianVariation * physicalColumn[scalarDof];
                        const double contribution = pressure * gradientWeightVariation;

                        actionColumn[scalarDof] -= contribution;
                    }
                }
            }

            MFEM_VERIFY(
                vector_is_finite(elementAction),
                "Prepared pressure-force displacement Jacobian encountered a non-finite element action."
            );

            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