module;

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

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;
    }
} // 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
    ) {
        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> displacementDofs;
        mfem::Array<int> compactificationDofs;

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

        mfem::Vector densityShape;
        mfem::Vector enthalpyShape;

        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.densityDofTransformation  = m_fem.densityFes->GetElementDofs(elementId, data.densityDofs);
            data.enthalpyDofTransformation = m_fem.enthalpyFes->GetElementDofs(elementId, data.enthalpyDofs);

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

            baseDensityLocal.GetSubVector(data.densityDofs, elementBaseDensity);
            baseEnthalpyLocal.GetSubVector(data.enthalpyDofs, elementBaseEnthalpy);
            displacementLocal.GetSubVector(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 (displacementDofTransformation != nullptr) {
                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.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::valid,
                    "Stateless mapping failed while preparing the barotropic closure operator. Element: "
                        << elementId << ", attribute: " << transformation->Attribute
                        << ", quadrature point: " << quadraturePoint << ", status: " << static_cast<int>(mappingStatus)
                );

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

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

                MFEM_VERIFY(
                    std::isfinite(quadratureWeight) && quadratureWeight > 0.0 && std::isfinite(eosDensity) &&
                        std::isfinite(enthalpyDerivative),
                    "PreparedBarotropicClosureOperator encountered invalid quadrature data."
                );

                data.quadratureWeights(quadraturePoint)          = quadratureWeight;
                data.weightedResidual(quadraturePoint)           = quadratureWeight * (density - eosDensity);
                data.weightedEnthalpyDerivative(quadraturePoint) = quadratureWeight * enthalpyDerivative;
            }
        }

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

        kernels::apply_barotropic_closure_displacement_action(
            m_fem, m_domainMapper, m_equationOfState, m_baseDensityTrue, m_baseEnthalpyTrue, m_baseDisplacementTrue,
            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);
    }

    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