module;

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

#include <mpi.h>

module mean_field;

import :operators.prepared_barotropic_closure;
import :operators.kernels.barotropic_closure;
import :field.registry;
import :utils.domain;

namespace {
    using DomainSchema  = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using ClosureDomain = mean_field::field::FieldDomainT<mean_field::field::Density>;

    void verify_required_spaces(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(f.mesh != nullptr, "PreparedBarotropicClosureOperator requires a mesh.");
        MFEM_VERIFY(
            f.densityFes != nullptr, "PreparedBarotropicClosureOperator requires the density finite-element space."
        );
        MFEM_VERIFY(
            f.enthalpyFes != nullptr, "PreparedBarotropicClosureOperator requires the enthalpy finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr,
            "PreparedBarotropicClosureOperator requires the displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr,
            "PreparedBarotropicClosureOperator requires the compactification finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr,
            "PreparedBarotropicClosureOperator requires the compactification coordinate."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "PreparedBarotropicClosureOperator requires the quadrature factory."
        );
    }

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

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

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

        localVector.SetSize(finiteElementSpace.GetVSize());

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

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

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

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

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

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

    [[nodiscard]] int get_eos_extra_order(const mean_field::eos::Polytrope &equationOfState) {
        const double extraOrder = (equationOfState.polytropic_index() - 1.0) *
                                  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 EOS effective polynomial order is invalid."
        );

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

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

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

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

        /*
         * The quadrature Query still carries the legacy DOMAINS metadata.
         * Element support itself is no longer selected through that enum;
         * support is determined above through Density::Support + DomainSchema.
         * The Query metadata can be migrated independently with the quadrature
         * subsystem without changing this operator's algebra.
         */
        const mean_field::quadrature::Query query =
            EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EosClosureSource>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(),
                std::array<int, 1>{get_eos_extra_order(equationOfState)}, mean_field::utils::DOMAINS::STELLAR,
                mean_field::quadrature::MappingKind::general
            );

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

        MFEM_VERIFY(
            resolution.integration_rule != nullptr,
            "The quadrature policy did not return a prepared EOS-closure integration rule."
        );

        return *resolution.integration_rule;
    }

    using BarotropicRejection       = mean_field::operators::BarotropicClosurePreparationRejection;
    using BarotropicRejectionReason = mean_field::operators::BarotropicClosurePreparationRejectionReason;

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

    /*
     * Rejections are selected by preparation phase, then by an explicit
     * detail priority.  An earlier phase wins: mapping, quadrature algebra,
     * then EOS evaluation.  Never depend on the declaration order or the
     * underlying integer representation of either public enum.
     */
    [[nodiscard]] int mapping_status_priority(const mean_field::mapping::MappingStatus status) {
        using Status = mean_field::mapping::MappingStatus;
        switch (status) {
        case Status::non_positive_determinant:
            return 7;
        case Status::non_finite_result:
            return 6;
        case Status::non_finite_input:
            return 5;
        case Status::outside_reference_domain:
            return 4;
        case Status::at_compactified_infinity:
            return 3;
        case Status::invalid_reference_radius:
            return 2;
        case Status::valid:
            throw std::logic_error("A valid mapping cannot be a barotropic candidate rejection.");
        case Status::invalid_dimension:
            throw std::logic_error("A mapping dimension error cannot be a barotropic candidate rejection.");
        }
        throw std::logic_error("Unknown mapping status in barotropic candidate rejection.");
    }

    [[nodiscard]] mean_field::mapping::MappingStatus mapping_status_from_priority(const int priority) {
        using Status = mean_field::mapping::MappingStatus;
        switch (priority) {
        case 7:
            return Status::non_positive_determinant;
        case 6:
            return Status::non_finite_result;
        case 5:
            return Status::non_finite_input;
        case 4:
            return Status::outside_reference_domain;
        case 3:
            return Status::at_compactified_infinity;
        case 2:
            return Status::invalid_reference_radius;
        default:
            throw std::logic_error("Invalid synchronized mapping priority for barotropic preparation.");
        }
    }

    [[nodiscard]] int eos_error_priority(const mean_field::eos::EvaluationErrorCode code) {
        using Code = mean_field::eos::EvaluationErrorCode;
        switch (code) {
        case Code::outside_domain:
            return 3;
        case Code::nonfinite_input:
            return 2;
        case Code::nonfinite_result:
            return 1;
        case Code::unsupported_relation:
        case Code::unsupported_derivative:
        case Code::wrong_input_count:
        case Code::wrong_input_quantity:
            throw std::logic_error("A structural EOS error cannot be a barotropic candidate rejection.");
        }
        throw std::logic_error("Unknown EOS error in barotropic candidate rejection.");
    }

    [[nodiscard]] mean_field::eos::EvaluationErrorCode eos_error_from_priority(const int priority) {
        using Code = mean_field::eos::EvaluationErrorCode;
        switch (priority) {
        case 3:
            return Code::outside_domain;
        case 2:
            return Code::nonfinite_input;
        case 1:
            return Code::nonfinite_result;
        default:
            throw std::logic_error("Invalid synchronized EOS priority for barotropic preparation.");
        }
    }

    [[nodiscard]] int rejection_priority(const BarotropicRejection &rejection) {
        switch (rejection.reason) {
        case BarotropicRejectionReason::mapping_failure:
            return 300 + mapping_status_priority(rejection.mappingStatus);
        case BarotropicRejectionReason::invalid_quadrature_data:
            return 200;
        case BarotropicRejectionReason::equation_of_state:
            return 100 + eos_error_priority(rejection.equationOfStateError);
        }
        throw std::logic_error("Unknown barotropic candidate-rejection reason.");
    }

    [[nodiscard]] BarotropicRejection rejection_from_priority(const int priority) {
        if (priority >= 300) {
            return {
                .reason        = BarotropicRejectionReason::mapping_failure,
                .mappingStatus = mapping_status_from_priority(priority - 300)
            };
        }
        if (priority == 200) {
            return {.reason = BarotropicRejectionReason::invalid_quadrature_data};
        }
        if (priority >= 100) {
            return {
                .reason               = BarotropicRejectionReason::equation_of_state,
                .equationOfStateError = eos_error_from_priority(priority - 100)
            };
        }
        throw std::logic_error("Invalid synchronized barotropic candidate-rejection priority.");
    }

    void retain_higher_priority_rejection(
        std::optional<BarotropicRejection> &current,
        const BarotropicRejection candidate
    ) {
        if (!current.has_value() || rejection_priority(candidate) > rejection_priority(*current)) {
            current = candidate;
        }
    }

    [[nodiscard]] std::optional<BarotropicRejection> synchronize_rejection(
        const std::optional<BarotropicRejection> &local,
        const MPI_Comm communicator
    ) {
        const int localPriority = local.has_value() ? rejection_priority(*local) : 0;
        int globalPriority      = 0;
        if (MPI_Allreduce(&localPriority, &globalPriority, 1, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS) {
            throw std::runtime_error("PreparedBarotropicClosureOperator could not synchronize candidate validity.");
        }
        if (globalPriority == 0) {
            return std::nullopt;
        }
        return rejection_from_priority(globalPriority);
    }
} // namespace

namespace mean_field::operators {
    struct PreparedBarotropicClosureOperator::ConstructionData final {
        field::FieldDofMap densityMap;
        field::FieldDofMap enthalpyMap;
        field::FieldDofMap displacementMap;

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

    PreparedBarotropicClosureOperator::ConstructionData
    PreparedBarotropicClosureOperator::MakeConstructionData(const fem::FEM &f) {
        verify_required_spaces(f);
        return ConstructionData(f);
    }

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

    PreparedBarotropicClosureOperator::PreparedBarotropicClosureOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &equationOfState,
        ConstructionData constructionData
    )
        : mfem::Operator(
              constructionData.densityMap.reduced_size(),
              constructionData.densityMap.reduced_size() + constructionData.enthalpyMap.reduced_size() +
                  constructionData.displacementMap.reduced_size()
          ),
          m_fem(f),
          m_domainMapper(domainMapper),
          m_equationOfState(equationOfState),
          m_densityMap(std::move(constructionData.densityMap)),
          m_enthalpyMap(std::move(constructionData.enthalpyMap)),
          m_displacementMap(std::move(constructionData.displacementMap)),
          m_context(
              f,
              domainMapper,
              m_densityMap,
              m_enthalpyMap,
              m_displacementMap
          ) {
        MFEM_VERIFY(
            m_densityMap.full_size() == m_fem.densityFes->GetTrueVSize(),
            "The density FieldDofMap does not match the density finite-element space."
        );
        MFEM_VERIFY(
            m_enthalpyMap.full_size() == m_fem.enthalpyFes->GetTrueVSize(),
            "The enthalpy FieldDofMap does not match the enthalpy finite-element space."
        );
        MFEM_VERIFY(
            m_displacementMap.full_size() == m_fem.displacementFes->GetTrueVSize(),
            "The displacement FieldDofMap does not match the displacement finite-element space."
        );

        m_baseDensityTrue.SetSize(m_densityMap.full_size());
        m_baseEnthalpyTrue.SetSize(m_enthalpyMap.full_size());
        m_baseDisplacementTrue.SetSize(m_displacementMap.full_size());

        m_densityVariationTrue.SetSize(m_densityMap.full_size());
        m_enthalpyVariationTrue.SetSize(m_enthalpyMap.full_size());
        m_displacementVariationTrue.SetSize(m_displacementMap.full_size());
        m_fullThermodynamicAction.SetSize(m_densityMap.full_size());
        m_fullDisplacementAction.SetSize(m_densityMap.full_size());
        m_fullResidual.SetSize(m_densityMap.full_size());

        m_baseDensityTrue           = 0.0;
        m_baseEnthalpyTrue          = 0.0;
        m_baseDisplacementTrue      = 0.0;
        m_densityVariationTrue      = 0.0;
        m_enthalpyVariationTrue     = 0.0;
        m_displacementVariationTrue = 0.0;
        m_fullThermodynamicAction   = 0.0;
        m_fullDisplacementAction    = 0.0;
        m_fullResidual              = 0.0;
    }

    PreparedBarotropicClosureReport PreparedBarotropicClosureOperator::Prepare(
        const context::barotropic::BarotropicClosureStateView &state,
        const context::barotropic::BarotropicClosureDependencies &dependencies
    ) {
        auto result = TryPrepare(state, dependencies);
        if (!result.has_value()) {
            const BarotropicClosurePreparationRejection &rejection = result.error();
            switch (rejection.reason) {
            case BarotropicClosurePreparationRejectionReason::equation_of_state:
                throw eos::EvaluationError(
                    rejection.equationOfStateError,
                    "PreparedBarotropicClosureOperator encountered invalid thermodynamic data."
                );
            case BarotropicClosurePreparationRejectionReason::invalid_quadrature_data:
                throw std::domain_error("PreparedBarotropicClosureOperator encountered non-finite quadrature data.");
            case BarotropicClosurePreparationRejectionReason::mapping_failure:
                throw std::domain_error("PreparedBarotropicClosureOperator could not map the candidate geometry.");
            }
            throw std::logic_error("Unknown barotropic candidate-rejection reason.");
        }
        return std::move(result).value();
    }

    BarotropicClosurePreparationResult PreparedBarotropicClosureOperator::TryPrepare(
        const context::barotropic::BarotropicClosureStateView &state,
        const context::barotropic::BarotropicClosureDependencies &dependencies
    ) {
        PreparedBarotropicClosureReport report;
        report.contextReport = m_context.Prepare(state, dependencies);

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

        /*
         * Canonical solver -> MFEM expansion. Unsupported density and
         * enthalpy DOFs are zero. Displacement is currently an identity map,
         * but it is deliberately routed through the same abstraction.
         */
        m_densityMap.scatter(m_context.GetBaseDensity(), m_baseDensityTrue);
        m_enthalpyMap.scatter(m_context.GetBaseEnthalpy(), m_baseEnthalpyTrue);
        m_displacementMap.scatter(m_context.GetDisplacement(), m_baseDisplacementTrue);

        m_isPrepared = false;
        m_elements.clear();
        m_elements.reserve(m_fem.mesh->GetNE());

        mfem::Vector baseDensityLocal;
        mfem::Vector baseEnthalpyLocal;
        mfem::Vector displacementLocal;

        true_to_local(*m_fem.densityFes, m_baseDensityTrue, baseDensityLocal);
        true_to_local(*m_fem.enthalpyFes, m_baseEnthalpyTrue, baseEnthalpyLocal);
        true_to_local(*m_fem.displacementFes, m_baseDisplacementTrue, displacementLocal);

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

        mfem::Array<int> compactificationDofs;

        mfem::Vector elementBaseDensity;
        mfem::Vector elementBaseEnthalpy;
        mfem::Vector elementDisplacement;
        mfem::Vector elementCompactification;

        mfem::Vector densityShape;
        mfem::Vector enthalpyShape;
        std::optional<BarotropicClosurePreparationRejection> localRejection;

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

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

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

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

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

            data.displacementDofTransformation =
                m_fem.displacementFes->GetElementVDofs(elementId, data.displacementDofs);
            mfem::DofTransformation *compactificationDofTransformation =
                m_fem.compactificationFes->GetElementDofs(elementId, compactificationDofs);

            baseDensityLocal.GetSubVector(data.densityDofs, elementBaseDensity);
            baseEnthalpyLocal.GetSubVector(data.enthalpyDofs, elementBaseEnthalpy);
            displacementLocal.GetSubVector(data.displacementDofs, elementDisplacement);
            m_fem.compactificationCoordinate->GetSubVector(compactificationDofs, elementCompactification);

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(elementBaseDensity);
            }
            if (data.enthalpyDofTransformation != nullptr) {
                data.enthalpyDofTransformation->InvTransformPrimal(elementBaseEnthalpy);
            }
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(elementDisplacement);
            }
            if (compactificationDofTransformation != nullptr) {
                compactificationDofTransformation->InvTransformPrimal(elementCompactification);
            }

            const mfem::FiniteElement &densityElement          = *m_fem.densityFes->GetFE(elementId);
            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);

            const mapping::ElementDisplacementData displacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementDisplacement);

            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, elementCompactification
            );

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

            const mfem::IntegrationRule &integrationRule =
                get_eos_rule(m_fem, m_equationOfState, densityElement, enthalpyElement, *transformation);

            const int quadraturePointCount = integrationRule.GetNPoints();
            const int densityDofCount      = densityElement.GetDof();
            const int enthalpyDofCount     = enthalpyElement.GetDof();

            data.densityBasis.SetSize(quadraturePointCount, densityDofCount);
            data.enthalpyBasis.SetSize(quadraturePointCount, enthalpyDofCount);
            data.inverseElementJacobians.SetSize(
                quadraturePointCount, m_fem.mesh->Dimension() * m_fem.mesh->Dimension()
            );
            data.weightedResidual.SetSize(quadraturePointCount);
            data.quadratureWeights.SetSize(quadraturePointCount);
            data.weightedEnthalpyDerivative.SetSize(quadraturePointCount);

            densityShape.SetSize(densityDofCount);
            enthalpyShape.SetSize(enthalpyDofCount);

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

                transformation->SetIntPoint(&integrationPoint);

                mapping::VolumeMappingContext mappingContext;

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

                MFEM_VERIFY(
                    mappingStatus != mapping::MappingStatus::invalid_dimension,
                    "Stateless mapping reported a dimension error while preparing the barotropic closure operator."
                );
                if (mappingStatus != mapping::MappingStatus::valid) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason        = BarotropicClosurePreparationRejectionReason::mapping_failure,
                                         .mappingStatus = mappingStatus}
                    );
                    continue;
                }

                MFEM_VERIFY(
                    !mappingContext.mapping.compactified,
                    "Prepared barotropic closure support unexpectedly includes a compactified element."
                );

                for (int row = 0; row < m_fem.mesh->Dimension(); ++row) {
                    for (int column = 0; column < m_fem.mesh->Dimension(); ++column) {
                        data.inverseElementJacobians(quadraturePoint, row * m_fem.mesh->Dimension() + column) =
                            mappingContext.quadrature.J_inv(row, column);
                    }
                }

                densityElement.CalcShape(integrationPoint, densityShape);
                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                if (!vector_is_finite(densityShape) || !vector_is_finite(enthalpyShape)) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason = BarotropicClosurePreparationRejectionReason::invalid_quadrature_data}
                    );
                    continue;
                }

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

                const double density          = elementBaseDensity * densityShape;
                const double enthalpy         = elementBaseEnthalpy * enthalpyShape;
                const double quadratureWeight = mappingContext.quadrature.weight;

                if (!std::isfinite(quadratureWeight) || quadratureWeight <= 0.0) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason = BarotropicClosurePreparationRejectionReason::invalid_quadrature_data}
                    );
                    continue;
                }

                if (!std::isfinite(density)) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason = BarotropicClosurePreparationRejectionReason::invalid_quadrature_data}
                    );
                    continue;
                }
                if (!std::isfinite(enthalpy)) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason = BarotropicClosurePreparationRejectionReason::equation_of_state,
                                         .equationOfStateError = eos::EvaluationErrorCode::nonfinite_input}
                    );
                    continue;
                }
                if (enthalpy < 0.0) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason = BarotropicClosurePreparationRejectionReason::equation_of_state,
                                         .equationOfStateError = eos::EvaluationErrorCode::outside_domain}
                    );
                    continue;
                }

                const dimensions::SpecificEnthalpyValue specificEnthalpy{enthalpy};
                const double eosDensity =
                    eos::evaluate<eos::quantity::Density>(m_equationOfState, specificEnthalpy).value();
                const double enthalpyDerivative =
                    eos::partialDerivative<eos::quantity::Density, eos::quantity::SpecificEnthalpy>(
                        m_equationOfState, specificEnthalpy
                    )
                        .value();

                if (!std::isfinite(eosDensity) || !std::isfinite(enthalpyDerivative)) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason = BarotropicClosurePreparationRejectionReason::equation_of_state,
                                         .equationOfStateError = eos::EvaluationErrorCode::nonfinite_result}
                    );
                    continue;
                }

                const double weightedResidual           = quadratureWeight * (density - eosDensity);
                const double weightedEnthalpyDerivative = quadratureWeight * enthalpyDerivative;
                if (!std::isfinite(weightedResidual) || !std::isfinite(weightedEnthalpyDerivative)) {
                    retain_higher_priority_rejection(
                        localRejection, {.reason = BarotropicClosurePreparationRejectionReason::invalid_quadrature_data}
                    );
                    continue;
                }

                data.quadratureWeights(quadraturePoint)          = quadratureWeight;
                data.weightedResidual(quadraturePoint)           = weightedResidual;
                data.weightedEnthalpyDerivative(quadraturePoint) = weightedEnthalpyDerivative;
            }
        }

        if (auto globalRejection = synchronize_rejection(localRejection, m_fem.densityFes->GetComm());
            globalRejection.has_value()) {
            return std::unexpected(*globalRejection);
        }

        MFEM_VERIFY(!m_elements.empty(), "PreparedBarotropicClosureOperator found no elements in Density::Support.");

        m_isPrepared = true;
        ++m_preparationCount;
        report.preparedElementData = true;
        return report;
    }

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

        mfem::Vector localResidual(m_fem.densityFes->GetVSize());
        localResidual = 0.0;

        mfem::Vector elementResidual;

        for (const ElementPAData &data : m_elements) {
            elementResidual.SetSize(data.densityDofs.Size());
            data.densityBasis.MultTranspose(data.weightedResidual, elementResidual);

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

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

        local_to_true(*m_fem.densityFes, localResidual, m_fullResidual);
        residual.SetSize(m_densityMap.reduced_size());
        m_densityMap.gather(m_fullResidual, residual);
    }

    void PreparedBarotropicClosureOperator::AssembleDensityJacobianDiagonal(mfem::Vector &diagonal) const {
        VerifyPrepared();

        mfem::Vector localDiagonal(m_fem.densityFes->GetVSize());
        localDiagonal = 0.0;
        mfem::Vector elementDiagonal;

        for (const ElementPAData &data : m_elements) {
            MFEM_VERIFY(
                data.densityDofTransformation == nullptr,
                "Density mass-diagonal assembly currently requires scalar L2 element DOFs without a DOF transform."
            );
            elementDiagonal.SetSize(data.densityDofs.Size());
            elementDiagonal = 0.0;
            for (int trialDof = 0; trialDof < data.densityDofs.Size(); ++trialDof) {
                for (int quadraturePoint = 0; quadraturePoint < data.quadratureWeights.Size(); ++quadraturePoint) {
                    const double basis = data.densityBasis(quadraturePoint, trialDof);
                    elementDiagonal(trialDof) += data.quadratureWeights(quadraturePoint) * basis * basis;
                }
            }
            localDiagonal.AddElementVector(data.densityDofs, elementDiagonal);
        }

        mfem::Vector trueDiagonal;
        local_to_true(*m_fem.densityFes, localDiagonal, trueDiagonal);
        diagonal.SetSize(m_densityMap.reduced_size());
        m_densityMap.gather(trueDiagonal, diagonal);
    }

    void PreparedBarotropicClosureOperator::Mult(
        const mfem::Vector &densityVariation,
        const mfem::Vector &enthalpyVariation,
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        MFEM_VERIFY(
            densityVariation.Size() == m_densityMap.reduced_size(),
            "The supported density-variation vector has the wrong size."
        );
        MFEM_VERIFY(
            enthalpyVariation.Size() == m_enthalpyMap.reduced_size(),
            "The supported enthalpy-variation vector has the wrong size."
        );
        MFEM_VERIFY(
            displacementVariation.Size() == m_displacementMap.reduced_size(),
            "The supported displacement-variation vector has the wrong size."
        );

        validate_finite_vector(densityVariation, "The density variation contains a non-finite value.");
        validate_finite_vector(enthalpyVariation, "The enthalpy variation contains a non-finite value.");
        validate_finite_vector(displacementVariation, "The displacement variation contains a non-finite value.");

        m_densityMap.scatter(densityVariation, m_densityVariationTrue);
        m_enthalpyMap.scatter(enthalpyVariation, m_enthalpyVariationTrue);
        m_displacementMap.scatter(displacementVariation, m_displacementVariationTrue);

        ApplyThermodynamicActionFull(m_densityVariationTrue, m_enthalpyVariationTrue, m_fullThermodynamicAction);

        ApplyDisplacementActionFull(m_displacementVariationTrue, m_fullDisplacementAction);

        MFEM_VERIFY(
            m_fullThermodynamicAction.Size() == m_densityMap.full_size() &&
                m_fullDisplacementAction.Size() == m_densityMap.full_size(),
            "A full barotropic-closure Jacobian action has an incompatible density-space size."
        );

        m_fullThermodynamicAction += m_fullDisplacementAction;
        action.SetSize(m_densityMap.reduced_size());
        m_densityMap.gather(m_fullThermodynamicAction, action);
    }

    void PreparedBarotropicClosureOperator::Mult(
        const mfem::Vector &combinedVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        MFEM_VERIFY(
            combinedVariation.Size() == Width(), "The packed supported barotropic-closure variation has the wrong size."
        );

        mfem::real_t *combinedData = const_cast<mfem::real_t *>(combinedVariation.HostRead());

        const int densitySize      = m_densityMap.reduced_size();
        const int enthalpySize     = m_enthalpyMap.reduced_size();
        const int displacementSize = m_displacementMap.reduced_size();

        const mfem::Vector densityVariation(combinedData, densitySize);
        const mfem::Vector enthalpyVariation(combinedData + densitySize, enthalpySize);
        const mfem::Vector displacementVariation(combinedData + densitySize + enthalpySize, displacementSize);

        Mult(densityVariation, enthalpyVariation, displacementVariation, action);
    }

    void PreparedBarotropicClosureOperator::ApplyThermodynamicActionFull(
        const mfem::Vector &densityVariationTrue,
        const mfem::Vector &enthalpyVariationTrue,
        mfem::Vector &actionTrue
    ) const {
        MFEM_VERIFY(
            densityVariationTrue.Size() == m_densityMap.full_size(), "The full density variation has the wrong size."
        );
        MFEM_VERIFY(
            enthalpyVariationTrue.Size() == m_enthalpyMap.full_size(), "The full enthalpy variation has the wrong size."
        );

        mfem::Vector densityVariationLocal;
        mfem::Vector enthalpyVariationLocal;

        true_to_local(*m_fem.densityFes, densityVariationTrue, densityVariationLocal);
        true_to_local(*m_fem.enthalpyFes, enthalpyVariationTrue, enthalpyVariationLocal);

        mfem::Vector localAction(m_fem.densityFes->GetVSize());
        localAction = 0.0;

        mfem::Vector elementDensityVariation;
        mfem::Vector elementEnthalpyVariation;
        mfem::Vector quadratureDensityVariation;
        mfem::Vector quadratureEnthalpyVariation;
        mfem::Vector quadratureAction;
        mfem::Vector elementAction;

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

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

            quadratureDensityVariation.SetSize(data.quadratureWeights.Size());
            quadratureEnthalpyVariation.SetSize(data.quadratureWeights.Size());
            quadratureAction.SetSize(data.quadratureWeights.Size());

            data.densityBasis.Mult(elementDensityVariation, quadratureDensityVariation);
            data.enthalpyBasis.Mult(elementEnthalpyVariation, quadratureEnthalpyVariation);

            for (int quadraturePoint = 0; quadraturePoint < quadratureAction.Size(); ++quadraturePoint) {
                quadratureAction(quadraturePoint) =
                    data.quadratureWeights(quadraturePoint) * quadratureDensityVariation(quadraturePoint) -
                    data.weightedEnthalpyDerivative(quadraturePoint) * quadratureEnthalpyVariation(quadraturePoint);
            }

            elementAction.SetSize(data.densityDofs.Size());
            data.densityBasis.MultTranspose(quadratureAction, elementAction);

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

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

        local_to_true(*m_fem.densityFes, localAction, actionTrue);
    }

    void PreparedBarotropicClosureOperator::ApplyDisplacementActionFull(
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &actionTrue
    ) const {
        MFEM_VERIFY(
            displacementVariationTrue.Size() == m_displacementMap.full_size(),
            "The full displacement variation has the wrong size."
        );

        true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacementVariationLocal);

        m_localDisplacementAction.SetSize(m_fem.densityFes->GetVSize());
        m_localDisplacementAction = 0.0;

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

        for (const ElementPAData &data : m_elements) {
            m_displacementVariationLocal.GetSubVector(data.displacementDofs, m_elementDisplacementVariation);
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(m_elementDisplacementVariation);
            }

            const mfem::FiniteElement &displacementElement = *m_fem.displacementFes->GetFE(data.elementId);
            const mapping::ElementDisplacementData directionData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, m_elementDisplacementVariation);
            const mfem::DenseMatrix &directionDofs      = directionData.GetDofMatrix();
            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);
            MFEM_VERIFY(
                transformation != nullptr,
                "Prepared barotropic closure displacement action received a null element transformation."
            );
            const mfem::FiniteElement &densityElement  = *m_fem.densityFes->GetFE(data.elementId);
            const mfem::FiniteElement &enthalpyElement = *m_fem.enthalpyFes->GetFE(data.elementId);
            const mfem::IntegrationRule &integrationRule =
                get_eos_rule(m_fem, m_equationOfState, densityElement, enthalpyElement, *transformation);

            MFEM_VERIFY(
                data.inverseElementJacobians.Height() == integrationRule.GetNPoints() &&
                    data.inverseElementJacobians.Width() == dimension * dimension,
                "Prepared barotropic closure inverse-Jacobian data has an incompatible size."
            );

            m_referenceDShape.SetSize(displacementElement.GetDof(), dimension);
            m_referenceDisplacementJacobian.SetSize(dimension, dimension);
            m_quadratureDisplacementAction.SetSize(integrationRule.GetNPoints());

            for (int quadraturePoint = 0; quadraturePoint < integrationRule.GetNPoints(); ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = integrationRule.IntPoint(quadraturePoint);
                displacementElement.CalcDShape(integrationPoint, m_referenceDShape);
                mfem::MultAtB(directionDofs, m_referenceDShape, m_referenceDisplacementJacobian);

                double logarithmicJacobianVariation{0.0};
                for (int row = 0; row < dimension; ++row) {
                    for (int column = 0; column < dimension; ++column) {
                        logarithmicJacobianVariation +=
                            data.inverseElementJacobians(quadraturePoint, row * dimension + column) *
                            m_referenceDisplacementJacobian(column, row);
                    }
                }

                m_quadratureDisplacementAction(quadraturePoint) =
                    data.weightedResidual(quadraturePoint) * logarithmicJacobianVariation;
                MFEM_VERIFY(
                    std::isfinite(m_quadratureDisplacementAction(quadraturePoint)),
                    "Prepared barotropic closure displacement action encountered a non-finite quadrature value."
                );
            }

            m_elementDisplacementAction.SetSize(data.densityDofs.Size());
            data.densityBasis.MultTranspose(m_quadratureDisplacementAction, m_elementDisplacementAction);

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->TransformDual(m_elementDisplacementAction);
            }
            m_localDisplacementAction.AddElementVector(data.densityDofs, m_elementDisplacementAction);
        }

        local_to_true(*m_fem.densityFes, m_localDisplacementAction, actionTrue);
    }

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

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

    int PreparedBarotropicClosureOperator::GetDensitySize() const noexcept {
        return m_densityMap.reduced_size();
    }

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

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

    const context::barotropic::BarotropicClosureLinearizationContext &
    PreparedBarotropicClosureOperator::GetContext() const noexcept {
        return m_context;
    }

    const context::barotropic::BarotropicClosurePreparationStatistics &
    PreparedBarotropicClosureOperator::GetContextPreparationStatistics() const noexcept {
        return m_context.GetPreparationStatistics();
    }

    void PreparedBarotropicClosureOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            m_isPrepared, "PreparedBarotropicClosureOperator must be prepared before this operation is called."
        );
    }
} // namespace mean_field::operators