module;

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>

#include <mfem.hpp>

module mean_field;

import :operators.prepared_hydrostatic_equilibrium;

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

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

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

    void copy_vector_block(
        const mfem::Vector &source,
        const int offset,
        const int size,
        mfem::Vector &block
    ) {
        MFEM_VERIFY(
            offset >= 0 && size >= 0 && offset + size <= source.Size(), "Hydrostatic Jacobian block lies outside the "
                                                                        "input vector."
        );

        block.SetSize(size);

        for (int entry = 0; entry < size; ++entry) {
            block(entry) = source(offset + entry);
        }
    }

    const mfem::IntegrationRule &get_hydrostatic_rule(
        const mean_field::fem::FEM &f,
        const mfem::FiniteElement &enthalpyElement,
        const mfem::FiniteElement &gravityPotentialElement,
        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 hydrostatic enthalpy element does "
            "not match the registered field."
        );

        MFEM_VERIFY(
            gravityPotentialElement.GetOrder() == mean_field::field::Gravity::Potential::familyOrder,
            "The prepared hydrostatic potential element does "
            "not match the registered field."
        );

        const auto enthalpyQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumEnthalpy>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        const auto gravityQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumGravity>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        // A rigid-rotation potential is quadratic in physical position.
        const auto rotationQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumRotation>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), std::array<int, 1>{2},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        const auto constantQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumConstant>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        const std::array<mean_field::quadrature::MfemRule, 4> candidateRules{
            f.quadratureFactory->get(enthalpyQuery, transformation.GetGeometryType()),
            f.quadratureFactory->get(gravityQuery, transformation.GetGeometryType()),
            f.quadratureFactory->get(rotationQuery, transformation.GetGeometryType()),
            f.quadratureFactory->get(constantQuery, transformation.GetGeometryType())
        };

        const auto selectedRule =
            std::max_element(candidateRules.begin(), candidateRules.end(), [](const auto &left, const auto &right) {
                return left.resolution.order < right.resolution.order;
            });

        MFEM_VERIFY(
            selectedRule != candidateRules.end() && selectedRule->integration_rule != nullptr,
            "The quadrature policy did not return a valid "
            "hydrostatic-equilibrium integration rule."
        );

        return *selectedRule->integration_rule;
    }
} // namespace

namespace mean_field::operators {
    HydrostaticJacobianBlockLayout::HydrostaticJacobianBlockLayout(const fem::FEM &f) {
        MFEM_VERIFY(
            f.enthalpyFes != nullptr, "HydrostaticJacobianBlockLayout requires the "
                                      "enthalpy finite-element space."
        );

        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "HydrostaticJacobianBlockLayout requires the "
                                              "gravity-potential finite-element space."
        );

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

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

        const field::FieldDofMap gravityPotentialMap =
            field::make_field_dof_map<field::Gravity, DomainSchema>(*f.gravityPotentialFes);

        const field::FieldDofMap displacementMap =
            field::make_field_dof_map<field::Displacement, DomainSchema>(*f.displacementFes);

        m_enthalpySize         = enthalpyMap.reduced_size();
        m_gravityPotentialSize = gravityPotentialMap.reduced_size();
        m_displacementSize     = displacementMap.reduced_size();
        m_residualSize         = m_enthalpySize;

        m_totalSize            = m_enthalpySize + m_gravityPotentialSize + 1 + m_displacementSize;

        MFEM_VERIFY(
            m_enthalpySize > 0 && m_gravityPotentialSize > 0 && m_displacementSize > 0,
            "HydrostaticJacobianBlockLayout received an empty "
            "finite-element space."
        );
    }

    int HydrostaticJacobianBlockLayout::Offset(const HydrostaticJacobianInputBlock block) const {
        switch (block) {
        case HydrostaticJacobianInputBlock::enthalpy:
            return 0;

        case HydrostaticJacobianInputBlock::gravityPotential:
            return m_enthalpySize;

        case HydrostaticJacobianInputBlock::bernoulliConstant:
            return m_enthalpySize + m_gravityPotentialSize;

        case HydrostaticJacobianInputBlock::displacement:
            return m_enthalpySize + m_gravityPotentialSize + 1;
        }

        MFEM_ABORT(
            "HydrostaticJacobianBlockLayout received an "
            "unknown input block."
        );

        return 0;
    }

    int HydrostaticJacobianBlockLayout::Size(const HydrostaticJacobianInputBlock block) const {
        switch (block) {
        case HydrostaticJacobianInputBlock::enthalpy:
            return m_enthalpySize;

        case HydrostaticJacobianInputBlock::gravityPotential:
            return m_gravityPotentialSize;

        case HydrostaticJacobianInputBlock::bernoulliConstant:
            return 1;

        case HydrostaticJacobianInputBlock::displacement:
            return m_displacementSize;
        }

        MFEM_ABORT(
            "HydrostaticJacobianBlockLayout received an "
            "unknown input block."
        );

        return 0;
    }

    int HydrostaticJacobianBlockLayout::GetTotalSize() const noexcept {
        return m_totalSize;
    }

    int HydrostaticJacobianBlockLayout::GetResidualSize() const noexcept {
        return m_residualSize;
    }

    PreparedHydrostaticEquilibriumOperator::PreparedHydrostaticEquilibriumOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper
    )
        : m_fem(f),
          m_domainMapper(domainMapper),
          m_context(
              f,
              domainMapper
          ) {
        MFEM_VERIFY(m_fem.mesh != nullptr, "PreparedHydrostaticEquilibriumOperator requires a mesh.");

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

        MFEM_VERIFY(
            m_fem.gravityPotentialFes != nullptr, "PreparedHydrostaticEquilibriumOperator requires "
                                                  "the gravity-potential finite-element space."
        );

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

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

        MFEM_VERIFY(
            m_fem.compactificationCoordinate != nullptr, "PreparedHydrostaticEquilibriumOperator requires "
                                                         "the compactification coordinate."
        );

        MFEM_VERIFY(
            m_fem.quadratureFactory != nullptr, "PreparedHydrostaticEquilibriumOperator requires "
                                                "the quadrature-rule factory."
        );

        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_fem.mesh->Dimension(), "The hydrostatic operator's stateless mapper "
                                                                      "dimension does not match the mesh dimension."
        );

        m_enthalpyVariationTrue.SetSize(m_context.GetEnthalpyMap().full_size());
        m_gravityPotentialVariationTrue.SetSize(m_context.GetGravityPotentialMap().full_size());
        m_displacementVariationTrue.SetSize(m_context.GetDisplacementMap().full_size());
        m_fullEnthalpyAction.SetSize(m_context.GetEnthalpyMap().full_size());
    }

    PreparedHydrostaticEquilibriumReport PreparedHydrostaticEquilibriumOperator::Prepare(
        const context::hydrostatic::HydrostaticEquilibriumStateView &state,
        const context::hydrostatic::HydrostaticEquilibriumDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        const bool rotationObjectChanged =
            !m_context.IsPrepared() || dependencies.rotation != m_context.GetDependencies().rotation;

        PreparedHydrostaticEquilibriumReport report;

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

        if (rotationObjectChanged) {
            m_rotation             = rotation;
            report.updatedRotation = true;
        }

        MFEM_VERIFY(
            m_rotation.has_value(), "The prepared hydrostatic operator has no frozen "
                                    "rotation state."
        );

        m_isPrepared = false;

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

        if (report.contextReport.preparedGeometryState) {
            PrepareGeometry();
            PrepareAlgebraicJacobianBlocks();
            report.preparedAlgebraicJacobianBlocks = true;
        }

        if (report.contextReport.preparedRotationDependencies) {
            PrepareRotation();
        }

        if (report.contextReport.preparedBaseState) {
            PrepareBaseState();
            FinalizeDisplacementJacobianPreparation();
            AssembleCachedResidual();
            ++m_residualPreparationCount;
            report.preparedDisplacementJacobianData = true;
            report.preparedResidual                 = true;
        }

        MFEM_VERIFY(
            !m_elements.empty(), "PreparedHydrostaticEquilibriumOperator found no "
                                 "stellar elements."
        );

        MFEM_VERIFY(
            m_cachedResidual.Size() == m_context.GetEnthalpyMap().reduced_size(),
            "The prepared hydrostatic residual has the wrong supported size."
        );

        m_isPrepared = true;
        return report;
    }

    void PreparedHydrostaticEquilibriumOperator::PrepareStaticPlan() {
        m_elements.clear();
        m_elements.reserve(m_fem.mesh->GetNE());

        mfem::Vector enthalpyShape;
        mfem::Vector gravityPotentialShape;

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

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

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

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

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

            MFEM_VERIFY(
                enthalpyElement.GetGeomType() == gravityPotentialElement.GetGeomType() &&
                    enthalpyElement.GetGeomType() == transformation->GetGeometryType(),
                "Hydrostatic 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.gravityPotentialDofTransformation =
                m_fem.gravityPotentialFes->GetElementDofs(elementId, data.gravityPotentialDofs);

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

            data.integrationRule =
                &get_hydrostatic_rule(m_fem, enthalpyElement, gravityPotentialElement, *transformation);

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

            const int enthalpyDofCount         = enthalpyElement.GetDof();

            const int gravityPotentialDofCount = gravityPotentialElement.GetDof();

            data.enthalpyBasis.SetSize(quadraturePointCount, enthalpyDofCount);

            data.gravityPotentialBasis.SetSize(quadraturePointCount, gravityPotentialDofCount);

            enthalpyShape.SetSize(enthalpyDofCount);
            gravityPotentialShape.SetSize(gravityPotentialDofCount);

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

                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                gravityPotentialElement.CalcShape(integrationPoint, gravityPotentialShape);

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

                for (int dof = 0; dof < gravityPotentialDofCount; ++dof) {
                    data.gravityPotentialBasis(quadraturePoint, dof) = gravityPotentialShape(dof);
                }
            }
        }
    }

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

        true_to_local(*m_fem.displacementFes, m_context.GetDisplacementTrue(), displacementLocal);

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

        mfem::Array<int> compactificationDofs;

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

        for (ElementPAData &data : m_elements) {
            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);

            MFEM_VERIFY(
                transformation != nullptr && data.integrationRule != nullptr,
                "Prepared hydrostatic geometry has invalid "
                "static element data."
            );

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

            displacementLocal.GetSubVector(data.displacementDofs, elementDisplacement);

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

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

            if (compactificationDofTransformation != nullptr) {
                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();

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

            data.quadratureWeights.SetSize(quadraturePointCount);

            data.baseMappingContexts.resize(quadraturePointCount);

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

                transformation->SetIntPoint(&integrationPoint);

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

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

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

                const double quadratureWeight = mappingContext.quadrature.weight;

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

                data.quadratureWeights(quadraturePoint) = quadratureWeight;

                for (int component = 0; component < m_fem.mesh->Dimension(); ++component) {
                    const double position = mappingContext.mapping.physical_position(component);

                    MFEM_VERIFY(
                        std::isfinite(position), "Prepared hydrostatic geometry encountered "
                                                 "a non-finite physical position."
                    );

                    data.physicalPositions(quadraturePoint, component) = position;
                }
            }
        }
    }

    void PreparedHydrostaticEquilibriumOperator::PrepareAlgebraicJacobianBlocks() {
        for (ElementPAData &data : m_elements) {
            const int quadraturePointCount     = data.quadratureWeights.Size();

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

            const int gravityPotentialDofCount = data.gravityPotentialBasis.Width();

            MFEM_VERIFY(
                data.enthalpyBasis.Height() == quadraturePointCount &&
                    data.gravityPotentialBasis.Height() == quadraturePointCount,
                "Prepared hydrostatic algebraic Jacobian has "
                "inconsistent quadrature data."
            );

            data.enthalpyJacobian.SetSize(enthalpyDofCount, enthalpyDofCount);

            data.gravityPotentialJacobian.SetSize(enthalpyDofCount, gravityPotentialDofCount);

            data.bernoulliConstantJacobian.SetSize(enthalpyDofCount);

            data.enthalpyJacobian          = 0.0;
            data.gravityPotentialJacobian  = 0.0;
            data.bernoulliConstantJacobian = 0.0;

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const double quadratureWeight = data.quadratureWeights(quadraturePoint);

                for (int testDof = 0; testDof < enthalpyDofCount; ++testDof) {
                    const double weightedTestBasis = quadratureWeight * data.enthalpyBasis(quadraturePoint, testDof);

                    data.bernoulliConstantJacobian(testDof) -= weightedTestBasis;

                    for (int trialDof = 0; trialDof < enthalpyDofCount; ++trialDof) {
                        data.enthalpyJacobian(testDof, trialDof) +=
                            weightedTestBasis * data.enthalpyBasis(quadraturePoint, trialDof);
                    }

                    for (int trialDof = 0; trialDof < gravityPotentialDofCount; ++trialDof) {
                        data.gravityPotentialJacobian(testDof, trialDof) +=
                            weightedTestBasis * data.gravityPotentialBasis(quadraturePoint, trialDof);
                    }
                }
            }
        }

        ++m_algebraicJacobianStatistics.preparations;
    }

    void PreparedHydrostaticEquilibriumOperator::PrepareRotation() {
        MFEM_VERIFY(m_rotation.has_value(), "Prepared hydrostatic rotation has no frozen state.");

        mfem::Vector physicalPosition(m_fem.mesh->Dimension());

        mfem::Vector coordinateDirection(m_fem.mesh->Dimension());

        for (ElementPAData &data : m_elements) {
            const int quadraturePointCount = data.physicalPositions.Height();

            MFEM_VERIFY(
                data.physicalPositions.Width() == m_fem.mesh->Dimension(), "Prepared hydrostatic rotation has invalid "
                                                                           "geometry data."
            );

            data.rotationPotential.SetSize(quadraturePointCount);

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

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                for (int component = 0; component < physicalPosition.Size(); ++component) {
                    physicalPosition(component) = data.physicalPositions(quadraturePoint, component);
                }

                const double rotationPotential = m_rotation->potential(physicalPosition);

                MFEM_VERIFY(
                    std::isfinite(rotationPotential), "Prepared hydrostatic rotation encountered "
                                                      "a non-finite potential."
                );

                data.rotationPotential(quadraturePoint) = rotationPotential;

                for (int component = 0; component < physicalPosition.Size(); ++component) {
                    coordinateDirection            = 0.0;
                    coordinateDirection(component) = 1.0;

                    const double gradientComponent =
                        m_rotation->potential_directional_derivative(physicalPosition, coordinateDirection);

                    MFEM_VERIFY(
                        std::isfinite(gradientComponent), "Prepared hydrostatic rotation encountered "
                                                          "a non-finite potential gradient."
                    );

                    data.rotationGradient(quadraturePoint, component) = gradientComponent;
                }
            }
        }
    }

    void PreparedHydrostaticEquilibriumOperator::PrepareBaseState() {
        mfem::Vector enthalpyLocal;
        mfem::Vector gravityPotentialLocal;

        true_to_local(*m_fem.enthalpyFes, m_context.GetBaseEnthalpyTrue(), enthalpyLocal);

        true_to_local(*m_fem.gravityPotentialFes, m_context.GetBaseGravityPotentialTrue(), gravityPotentialLocal);

        mfem::Vector elementEnthalpy;
        mfem::Vector elementGravityPotential;
        mfem::Vector quadratureEnthalpy;
        mfem::Vector quadratureGravityPotential;

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

            gravityPotentialLocal.GetSubVector(data.gravityPotentialDofs, elementGravityPotential);

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

            if (data.gravityPotentialDofTransformation != nullptr) {
                data.gravityPotentialDofTransformation->InvTransformPrimal(elementGravityPotential);
            }

            const int quadraturePointCount = data.quadratureWeights.Size();

            quadratureEnthalpy.SetSize(quadraturePointCount);

            quadratureGravityPotential.SetSize(quadraturePointCount);

            data.enthalpyBasis.Mult(elementEnthalpy, quadratureEnthalpy);

            data.gravityPotentialBasis.Mult(elementGravityPotential, quadratureGravityPotential);

            MFEM_VERIFY(
                data.rotationPotential.Size() == quadraturePointCount, "Prepared hydrostatic base state has stale "
                                                                       "rotation data."
            );

            data.weightedResidual.SetSize(quadraturePointCount);

            data.hydrostaticImbalance.SetSize(quadraturePointCount);

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const double imbalance = quadratureEnthalpy(quadraturePoint) +
                                         quadratureGravityPotential(quadraturePoint) -
                                         data.rotationPotential(quadraturePoint) - m_context.GetBernoulliConstant();

                const double weightedResidual = data.quadratureWeights(quadraturePoint) * imbalance;

                MFEM_VERIFY(
                    std::isfinite(weightedResidual), "Prepared hydrostatic base state encountered "
                                                     "a non-finite residual value."
                );

                data.weightedResidual(quadraturePoint)     = weightedResidual;

                data.hydrostaticImbalance(quadraturePoint) = imbalance;
            }
        }
    }

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

        for (const ElementPAData &data : m_elements) {
            const int quadraturePointCount = data.quadratureWeights.Size();

            MFEM_VERIFY(
                data.baseDisplacementData.has_value() && data.compactificationData.has_value() &&
                    static_cast<int>(data.baseMappingContexts.size()) == quadraturePointCount &&
                    data.rotationGradient.Height() == quadraturePointCount &&
                    data.rotationGradient.Width() == dimension &&
                    data.hydrostaticImbalance.Size() == quadraturePointCount,
                "Prepared hydrostatic displacement Jacobian "
                "has inconsistent frozen data."
            );
        }

        ++m_displacementJacobianStatistics.preparations;
    }

    void PreparedHydrostaticEquilibriumOperator::AssembleCachedResidual() {
        mfem::Vector localResidual(m_fem.enthalpyFes->GetVSize());

        localResidual = 0.0;
        mfem::Vector elementResidual;

        for (const ElementPAData &data : m_elements) {
            elementResidual.SetSize(data.enthalpyDofs.Size());

            data.enthalpyBasis.MultTranspose(data.weightedResidual, elementResidual);

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

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

        local_to_true(*m_fem.enthalpyFes, localResidual, m_fullEnthalpyAction);

        m_cachedResidual.SetSize(m_context.GetEnthalpyMap().reduced_size());
        m_context.GetEnthalpyMap().gather(m_fullEnthalpyAction, m_cachedResidual);
    }

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

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

        MFEM_VERIFY(
            enthalpyVariation.Size() == m_context.GetEnthalpyMap().reduced_size(),
            "Prepared hydrostatic enthalpy variation has the wrong supported size."
        );

        m_context.GetEnthalpyMap().scatter(enthalpyVariation, m_enthalpyVariationTrue);

        mfem::Vector enthalpyVariationLocal;

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

        mfem::Vector localAction(m_fem.enthalpyFes->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.enthalpyDofTransformation != nullptr) {
                data.enthalpyDofTransformation->TransformDual(elementAction);
            }

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

        local_to_true(*m_fem.enthalpyFes, localAction, m_fullEnthalpyAction);

        action.SetSize(m_context.GetEnthalpyMap().reduced_size());
        m_context.GetEnthalpyMap().gather(m_fullEnthalpyAction, action);

        ++m_algebraicJacobianStatistics.enthalpyApplications;
    }

    void PreparedHydrostaticEquilibriumOperator::AssembleEnthalpyJacobianDiagonal(mfem::Vector &diagonal) const {
        VerifyPrepared();

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

        for (const ElementPAData &data : m_elements) {
            MFEM_VERIFY(
                data.enthalpyDofTransformation == nullptr,
                "Enthalpy mass-diagonal assembly currently requires scalar H1 element DOFs without a DOF transform."
            );
            MFEM_VERIFY(
                data.enthalpyJacobian.Height() == data.enthalpyDofs.Size() &&
                    data.enthalpyJacobian.Width() == data.enthalpyDofs.Size(),
                "The prepared enthalpy Jacobian block is not square on an element."
            );
            elementDiagonal.SetSize(data.enthalpyDofs.Size());
            for (int dof = 0; dof < data.enthalpyDofs.Size(); ++dof) {
                elementDiagonal(dof) = data.enthalpyJacobian(dof, dof);
            }
            localDiagonal.AddElementVector(data.enthalpyDofs, elementDiagonal);
        }

        mfem::Vector trueDiagonal;
        local_to_true(*m_fem.enthalpyFes, localDiagonal, trueDiagonal);
        diagonal.SetSize(m_context.GetEnthalpyMap().reduced_size());
        m_context.GetEnthalpyMap().gather(trueDiagonal, diagonal);
    }

    void PreparedHydrostaticEquilibriumOperator::ApplyGravityPotentialJacobianAction(
        const mfem::Vector &gravityPotentialVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        MFEM_VERIFY(
            gravityPotentialVariation.Size() == m_context.GetGravityPotentialMap().reduced_size(),
            "Prepared hydrostatic gravity-potential variation has the wrong "
            "supported size."
        );

        m_context.GetGravityPotentialMap().scatter(gravityPotentialVariation, m_gravityPotentialVariationTrue);

        mfem::Vector gravityPotentialVariationLocal;

        true_to_local(*m_fem.gravityPotentialFes, m_gravityPotentialVariationTrue, gravityPotentialVariationLocal);

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

        localAction = 0.0;

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

        for (const ElementPAData &data : m_elements) {
            gravityPotentialVariationLocal.GetSubVector(data.gravityPotentialDofs, elementVariation);

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

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

            data.gravityPotentialJacobian.Mult(elementVariation, elementAction);

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

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

        local_to_true(*m_fem.enthalpyFes, localAction, m_fullEnthalpyAction);

        action.SetSize(m_context.GetEnthalpyMap().reduced_size());
        m_context.GetEnthalpyMap().gather(m_fullEnthalpyAction, action);

        ++m_algebraicJacobianStatistics.gravityPotentialApplications;
    }

    void PreparedHydrostaticEquilibriumOperator::ApplyBernoulliConstantJacobianAction(
        const double bernoulliConstantVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        MFEM_VERIFY(
            std::isfinite(bernoulliConstantVariation), "Prepared hydrostatic Bernoulli-constant Jacobian "
                                                       "received a non-finite variation."
        );

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

        localAction = 0.0;
        mfem::Vector elementAction;

        for (const ElementPAData &data : m_elements) {
            elementAction = data.bernoulliConstantJacobian;
            elementAction *= bernoulliConstantVariation;

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

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

        local_to_true(*m_fem.enthalpyFes, localAction, m_fullEnthalpyAction);

        action.SetSize(m_context.GetEnthalpyMap().reduced_size());
        m_context.GetEnthalpyMap().gather(m_fullEnthalpyAction, action);

        ++m_algebraicJacobianStatistics.bernoulliConstantApplications;
    }

    void PreparedHydrostaticEquilibriumOperator::ApplyAlgebraicJacobianAction(
        const mfem::Vector &enthalpyVariation,
        const mfem::Vector &gravityPotentialVariation,
        const double bernoulliConstantVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        MFEM_VERIFY(
            std::isfinite(bernoulliConstantVariation), "Prepared hydrostatic algebraic Jacobian received "
                                                       "a non-finite Bernoulli-constant variation."
        );

        MFEM_VERIFY(
            enthalpyVariation.Size() == m_context.GetEnthalpyMap().reduced_size(),
            "Prepared hydrostatic algebraic enthalpy variation has the wrong "
            "supported size."
        );

        MFEM_VERIFY(
            gravityPotentialVariation.Size() == m_context.GetGravityPotentialMap().reduced_size(),
            "Prepared hydrostatic algebraic gravity-potential variation has "
            "the wrong supported size."
        );

        m_context.GetEnthalpyMap().scatter(enthalpyVariation, m_enthalpyVariationTrue);
        m_context.GetGravityPotentialMap().scatter(gravityPotentialVariation, m_gravityPotentialVariationTrue);

        mfem::Vector enthalpyVariationLocal;
        mfem::Vector gravityPotentialVariationLocal;

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

        true_to_local(*m_fem.gravityPotentialFes, m_gravityPotentialVariationTrue, gravityPotentialVariationLocal);

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

        localAction = 0.0;

        mfem::Vector elementEnthalpyVariation;
        mfem::Vector elementGravityPotentialVariation;
        mfem::Vector elementAction;
        mfem::Vector elementWorkspace;

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

            gravityPotentialVariationLocal.GetSubVector(data.gravityPotentialDofs, elementGravityPotentialVariation);

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

            if (data.gravityPotentialDofTransformation != nullptr) {
                data.gravityPotentialDofTransformation->InvTransformPrimal(elementGravityPotentialVariation);
            }

            MFEM_VERIFY(
                data.enthalpyJacobian.Height() == data.gravityPotentialJacobian.Height() &&
                    data.enthalpyJacobian.Width() == elementEnthalpyVariation.Size() &&
                    data.gravityPotentialJacobian.Width() == elementGravityPotentialVariation.Size() &&
                    data.bernoulliConstantJacobian.Size() == data.enthalpyJacobian.Height(),
                "Prepared hydrostatic algebraic Jacobian has "
                "incompatible element dimensions."
            );

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

            elementWorkspace.SetSize(data.gravityPotentialJacobian.Height());

            data.enthalpyJacobian.Mult(elementEnthalpyVariation, elementAction);

            data.gravityPotentialJacobian.Mult(elementGravityPotentialVariation, elementWorkspace);

            elementAction.Add(1.0, elementWorkspace);
            elementAction.Add(bernoulliConstantVariation, data.bernoulliConstantJacobian);

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

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

        local_to_true(*m_fem.enthalpyFes, localAction, m_fullEnthalpyAction);

        action.SetSize(m_context.GetEnthalpyMap().reduced_size());
        m_context.GetEnthalpyMap().gather(m_fullEnthalpyAction, action);

        ++m_algebraicJacobianStatistics.combinedApplications;
    }

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

        MFEM_VERIFY(
            displacementVariation.Size() == m_context.GetDisplacementMap().reduced_size(),
            "Prepared hydrostatic displacement variation has the wrong "
            "supported size."
        );

        m_context.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        mfem::Vector displacementVariationLocal;

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

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

        localAction = 0.0;

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

        mfem::Vector elementDisplacementVariation;
        mfem::Vector weightedQuadratureVariation;
        mfem::Vector elementAction;
        mapping::VolumeMappingVariation variation;

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

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

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

            displacementVariationLocal.GetSubVector(data.displacementDofs, elementDisplacementVariation);

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

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

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

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

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

            MFEM_VERIFY(
                static_cast<int>(data.baseMappingContexts.size()) == quadraturePointCount &&
                    data.quadratureWeights.Size() == quadraturePointCount &&
                    data.hydrostaticImbalance.Size() == quadraturePointCount &&
                    data.rotationGradient.Height() == quadraturePointCount &&
                    data.rotationGradient.Width() == m_fem.mesh->Dimension(),
                "Prepared hydrostatic displacement Jacobian "
                "has inconsistent quadrature data."
            );

            weightedQuadratureVariation.SetSize(quadraturePointCount);

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

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

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

                double rotationPotentialVariation = 0.0;

                for (int component = 0; component < m_fem.mesh->Dimension(); ++component) {
                    rotationPotentialVariation += data.rotationGradient(quadraturePoint, component) *
                                                  variation.mapping.physical_position_variation(component);
                }

                const double weightedVariation =
                    data.hydrostaticImbalance(quadraturePoint) * variation.weight_variation -
                    data.quadratureWeights(quadraturePoint) * rotationPotentialVariation;

                MFEM_VERIFY(
                    std::isfinite(weightedVariation), "Prepared hydrostatic displacement Jacobian "
                                                      "encountered a non-finite quadrature action."
                );

                weightedQuadratureVariation(quadraturePoint) = weightedVariation;
            }

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

            data.enthalpyBasis.MultTranspose(weightedQuadratureVariation, elementAction);

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

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

        local_to_true(*m_fem.enthalpyFes, localAction, m_fullEnthalpyAction);

        action.SetSize(m_context.GetEnthalpyMap().reduced_size());
        m_context.GetEnthalpyMap().gather(m_fullEnthalpyAction, action);

        ++m_displacementJacobianStatistics.applications;
    }

    void PreparedHydrostaticEquilibriumOperator::ApplyCompleteJacobianAction(
        const mfem::Vector &enthalpyVariation,
        const mfem::Vector &gravityPotentialVariation,
        const double bernoulliConstantVariation,
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        mfem::Vector displacementAction;

        ApplyAlgebraicJacobianAction(enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, action);

        ApplyDisplacementJacobianAction(displacementVariation, displacementAction);

        MFEM_VERIFY(
            action.Size() == displacementAction.Size(), "Prepared hydrostatic complete Jacobian produced "
                                                        "incompatible algebraic and displacement actions."
        );

        action += displacementAction;
        ++m_completeJacobianStatistics.applications;
    }

    bool PreparedHydrostaticEquilibriumOperator::IsPrepared() const noexcept {
        return m_isPrepared;
    }

    const context::hydrostatic::HydrostaticPreparationStatistics &
    PreparedHydrostaticEquilibriumOperator::GetContextPreparationStatistics() const noexcept {
        return m_context.GetPreparationStatistics();
    }

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

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

    const PreparedHydrostaticAlgebraicJacobianStatistics &
    PreparedHydrostaticEquilibriumOperator::GetAlgebraicJacobianStatistics() const noexcept {
        return m_algebraicJacobianStatistics;
    }

    const PreparedHydrostaticDisplacementJacobianStatistics &
    PreparedHydrostaticEquilibriumOperator::GetDisplacementJacobianStatistics() const noexcept {
        return m_displacementJacobianStatistics;
    }

    const PreparedHydrostaticCompleteJacobianStatistics &
    PreparedHydrostaticEquilibriumOperator::GetCompleteJacobianStatistics() const noexcept {
        return m_completeJacobianStatistics;
    }

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

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

    const field::FieldDofMap &PreparedHydrostaticEquilibriumOperator::GetEnthalpyMap() const noexcept {
        return m_context.GetEnthalpyMap();
    }

    const field::FieldDofMap &PreparedHydrostaticEquilibriumOperator::GetGravityPotentialMap() const noexcept {
        return m_context.GetGravityPotentialMap();
    }

    const field::FieldDofMap &PreparedHydrostaticEquilibriumOperator::GetDisplacementMap() const noexcept {
        return m_context.GetDisplacementMap();
    }

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

    PreparedHydrostaticEquilibriumJacobianOperator::PreparedHydrostaticEquilibriumJacobianOperator(
        const fem::FEM &f,
        const PreparedHydrostaticEquilibriumOperator &preparedOperator
    )
        : mfem::Operator(
              HydrostaticJacobianBlockLayout(f).GetResidualSize(),
              HydrostaticJacobianBlockLayout(f).GetTotalSize()
          ),
          m_layout(f),
          m_preparedOperator(preparedOperator) {
        MFEM_VERIFY(
            &m_preparedOperator.GetFEM() == &f, "Prepared hydrostatic MFEM adapter and prepared "
                                                "operator must use the same FEM object."
        );

        MFEM_VERIFY(
            Height() == m_layout.GetResidualSize() && Width() == m_layout.GetTotalSize(),
            "Prepared hydrostatic MFEM adapter has "
            "inconsistent operator dimensions."
        );
    }

    void PreparedHydrostaticEquilibriumJacobianOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            direction.Size() == Width(), "Prepared hydrostatic MFEM adapter received a "
                                         "direction with the wrong size."
        );

        mfem::Vector enthalpyVariation;
        mfem::Vector gravityPotentialVariation;
        mfem::Vector displacementVariation;

        copy_vector_block(
            direction, m_layout.Offset(HydrostaticJacobianInputBlock::enthalpy),
            m_layout.Size(HydrostaticJacobianInputBlock::enthalpy), enthalpyVariation
        );

        copy_vector_block(
            direction, m_layout.Offset(HydrostaticJacobianInputBlock::gravityPotential),
            m_layout.Size(HydrostaticJacobianInputBlock::gravityPotential), gravityPotentialVariation
        );

        copy_vector_block(
            direction, m_layout.Offset(HydrostaticJacobianInputBlock::displacement),
            m_layout.Size(HydrostaticJacobianInputBlock::displacement), displacementVariation
        );

        const double bernoulliConstantVariation =
            direction(m_layout.Offset(HydrostaticJacobianInputBlock::bernoulliConstant));

        m_preparedOperator.ApplyCompleteJacobianAction(
            enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, displacementVariation, action
        );

        MFEM_VERIFY(
            action.Size() == Height(), "Prepared hydrostatic MFEM adapter produced an "
                                       "action with the wrong size."
        );
    }

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