module;

#include <algorithm>
#include <array>
#include <cmath>
#include <expected>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>

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

module mean_field;

import :operators.prepared_angular_momentum;

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

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

    void verify_finite_vector(
        const mfem::Vector &vector,
        const char *message
    ) {
        MFEM_VERIFY(is_finite_vector(vector), message);
    }

    [[nodiscard]] bool is_candidate_mapping_failure(const mean_field::mapping::MappingStatus status) {
        using mean_field::mapping::MappingStatus;
        return status == MappingStatus::non_finite_input || status == MappingStatus::non_finite_result ||
               status == MappingStatus::non_positive_determinant;
    }

    [[nodiscard]] std::optional<mean_field::mapping::MappingStatus> synchronize_mapping_failure(
        const std::optional<mean_field::mapping::MappingStatus> localFailure,
        const MPI_Comm communicator
    ) {
        int localFailures[2]{0, 0};
        if (localFailure.has_value()) {
            const int encodedStatus = static_cast<int>(*localFailure) + 1;
            if (is_candidate_mapping_failure(*localFailure)) {
                localFailures[0] = encodedStatus;
            } else {
                localFailures[1] = encodedStatus;
            }
        }

        int globalFailures[2]{0, 0};
        if (MPI_Allreduce(localFailures, globalFailures, 2, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS) {
            throw std::runtime_error("PreparedAngularMomentumOperator could not synchronize mapped-geometry validity.");
        }
        if (globalFailures[1] != 0) {
            throw std::runtime_error(
                "PreparedAngularMomentumOperator encountered a structural mapping failure with status " +
                std::to_string(globalFailures[1] - 1) + "."
            );
        }
        if (globalFailures[0] == 0) {
            return std::nullopt;
        }
        return static_cast<mean_field::mapping::MappingStatus>(globalFailures[0] - 1);
    }

    [[nodiscard]] bool synchronize_non_finite_failure(
        const bool localFailure,
        const MPI_Comm communicator,
        const char *operation
    ) {
        const int localStatus = localFailure ? 1 : 0;
        int globalStatus      = 0;
        if (MPI_Allreduce(&localStatus, &globalStatus, 1, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS) {
            throw std::runtime_error(
                std::string("PreparedAngularMomentumOperator could not synchronize ") + operation + "."
            );
        }
        return globalStatus != 0;
    }

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

    [[nodiscard]] const mfem::IntegrationRule &get_moment_of_inertia_rule(
        const mean_field::fem::FEM &f,
        const mfem::FiniteElement &densityElement,
        const mfem::ElementTransformation &transformation
    ) {
        using DensityField = mean_field::field::Field<mean_field::field::Density>;
        MFEM_VERIFY(
            densityElement.GetOrder() == mean_field::field::Density::Scalar::familyOrder,
            "The angular-momentum element does not match the registered density field."
        );
        const mean_field::quadrature::Query query =
            DensityField::make_query<mean_field::field::Density::Form::Quadrupole>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), std::array<int, 1>{2},
                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 an angular-momentum integration rule."
        );
        return *resolution.integration_rule;
    }

    void validate_shared_gravity_revisions(
        const mean_field::operators::context::gravity_field::GravityFieldLinearizationContext &gravityContext,
        const mean_field::operators::AngularMomentumDependencies &dependencies
    ) {
        MFEM_VERIFY(
            gravityContext.IsPrepared(),
            "PreparedAngularMomentumOperator requires the shared gravity context to be prepared first."
        );
        const auto &revisions = gravityContext.GetRevisions();
        MFEM_VERIFY(
            revisions.discretization.value == dependencies.discretization.revision &&
                revisions.density.value == dependencies.density.revision &&
                revisions.displacement.value == dependencies.displacement.revision,
            "PreparedAngularMomentumOperator received revisions that do not match the shared gravity context."
        );
    }

    void validate_identity_transition(
        const mean_field::operators::AngularMomentumDependencyStamp &prepared,
        const mean_field::operators::AngularMomentumDependencyStamp &requested,
        const char *message
    ) {
        MFEM_VERIFY(prepared.identity == requested.identity || prepared.revision != requested.revision, message);
    }
} // namespace

namespace mean_field::operators {
    PreparedAngularMomentumOperator::PreparedAngularMomentumOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const context::gravity_field::GravityFieldLinearizationContext &gravityContext,
        models::CompiledFixedAngularMomentum constraint
    )
        : m_fem(f),
          m_domainMapper(domainMapper),
          m_gravityContext(gravityContext),
          m_constraint(std::move(constraint)) {
        MFEM_VERIFY(m_fem.mesh != nullptr, "PreparedAngularMomentumOperator requires a mesh.");
        MFEM_VERIFY(
            m_fem.mesh->Dimension() == 3 && m_domainMapper.GetDimension() == 3,
            "PreparedAngularMomentumOperator currently requires a three-dimensional mapped domain."
        );
        MFEM_VERIFY(
            m_fem.densityFes != nullptr && m_fem.displacementFes != nullptr && m_fem.compactificationFes != nullptr &&
                m_fem.compactificationCoordinate != nullptr && m_fem.quadratureFactory != nullptr,
            "PreparedAngularMomentumOperator requires density, displacement, compactification, and quadrature data."
        );
        MFEM_VERIFY(
            m_gravityContext.GetDensityMap().full_size() == m_fem.densityFes->GetTrueVSize() &&
                m_gravityContext.GetDisplacementMap().full_size() == m_fem.displacementFes->GetTrueVSize(),
            "PreparedAngularMomentumOperator received incompatible shared FieldDof maps."
        );
        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
    }

    PreparedAngularMomentumReport PreparedAngularMomentumOperator::Prepare(
        const double angularVelocity,
        const AngularMomentumDependencies &dependencies
    ) {
        auto result = TryPrepare(angularVelocity, dependencies);
        if (!result.has_value()) {
            throwAngularMomentumPreparationRejection(result.error());
        }
        return std::move(result).value();
    }

    AngularMomentumPreparationResult PreparedAngularMomentumOperator::TryPrepare(
        const double angularVelocity,
        const AngularMomentumDependencies &dependencies
    ) {
        validate_shared_gravity_revisions(m_gravityContext, dependencies);

        if (m_isPrepared) {
            validate_identity_transition(
                m_preparedDependencies.discretization, dependencies.discretization,
                "A new angular-momentum discretization identity must change its revision."
            );
            validate_identity_transition(
                m_preparedDependencies.density, dependencies.density,
                "A new angular-momentum density identity must change its revision."
            );
            validate_identity_transition(
                m_preparedDependencies.displacement, dependencies.displacement,
                "A new angular-momentum displacement identity must change its revision."
            );
            validate_identity_transition(
                m_preparedDependencies.rotation, dependencies.rotation,
                "A new angular-momentum rotation identity must change its revision."
            );
        }

        const bool rebuildStaticPlan =
            !m_isPrepared || dependencies.discretization != m_preparedDependencies.discretization;
        const bool refreshGeometry =
            rebuildStaticPlan || dependencies.displacement != m_preparedDependencies.displacement;
        const bool refreshDensity        = rebuildStaticPlan || dependencies.density != m_preparedDependencies.density;
        const bool updateAngularVelocity = !m_isPrepared || dependencies.rotation != m_preparedDependencies.rotation ||
                                           angularVelocity != m_angularVelocity;

        m_isPrepared                     = false;
        if (synchronize_non_finite_failure(
                !std::isfinite(angularVelocity), m_fem.mesh->GetComm(), "angular-velocity validity"
            )) {
            return std::unexpected(
                AngularMomentumPreparationRejection{
                    .reason = AngularMomentumPreparationRejectionReason::non_finite_angular_velocity
                }
            );
        }

        PreparedAngularMomentumReport report;
        if (rebuildStaticPlan) {
            BuildStaticPlan();
            report.rebuiltStaticPlan = true;
        }
        if (refreshGeometry) {
            const auto mappingFailure = synchronize_mapping_failure(
                RefreshGeometry(m_gravityContext.GetGeometryContext().GetDisplacementTrue()), m_fem.mesh->GetComm()
            );
            if (mappingFailure.has_value()) {
                const auto reason = *mappingFailure == mapping::MappingStatus::non_positive_determinant
                                        ? AngularMomentumPreparationRejectionReason::inverted_geometry
                                        : AngularMomentumPreparationRejectionReason::non_finite_geometry;
                return std::unexpected(
                    AngularMomentumPreparationRejection{.reason = reason, .mappingStatus = *mappingFailure}
                );
            }
            report.refreshedGeometry = true;
        }
        if (refreshDensity) {
            if (synchronize_non_finite_failure(
                    !RefreshDensity(m_gravityContext.GetDensityTrue()), m_fem.mesh->GetComm(),
                    "interpolated-density validity"
                )) {
                return std::unexpected(
                    AngularMomentumPreparationRejection{
                        .reason = AngularMomentumPreparationRejectionReason::non_finite_density
                    }
                );
            }
            report.refreshedDensity = true;
        }
        if (updateAngularVelocity) {
            m_angularVelocity             = angularVelocity;
            report.updatedAngularVelocity = true;
        }
        if (refreshGeometry || refreshDensity || updateAngularVelocity) {
            auto rejection = TryAssembleResidual();
            if (rejection.has_value()) {
                return std::unexpected(*rejection);
            }
            report.assembledResidual = true;
        }

        m_preparedDependencies = dependencies;
        m_isPrepared           = true;
        return report;
    }

    void PreparedAngularMomentumOperator::BuildStaticPlan() {
        m_elements.clear();
        m_elements.reserve(m_fem.mesh->GetNE());
        int localStellarElementCount = 0;
        for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(elementId);
            MFEM_VERIFY(transformation != nullptr, "Angular-momentum preparation received a null transformation.");
            if (is_vacuum_attribute(transformation->Attribute)) {
                continue;
            }
            ++localStellarElementCount;
            m_elements.emplace_back();
            ElementPAData &data           = m_elements.back();
            data.elementId                = elementId;
            data.densityDofTransformation = m_fem.densityFes->GetElementDofs(elementId, data.densityDofs);
            data.displacementDofTransformation =
                m_fem.displacementFes->GetElementVDofs(elementId, data.displacementDofs);
            data.compactificationDofTransformation =
                m_fem.compactificationFes->GetElementDofs(elementId, data.compactificationDofs);

            const mfem::FiniteElement &densityElement = *m_fem.densityFes->GetFE(elementId);
            const mfem::IntegrationRule &integrationRule =
                get_moment_of_inertia_rule(m_fem, densityElement, *transformation);
            data.integrationRule = &integrationRule;
            data.densityBasis    = m_fem.GetReferenceTables().GetScalarTable(densityElement, integrationRule);
            data.mappingContexts.SetSize(integrationRule.GetNPoints(), m_fem.mesh->Dimension());
            data.density.SetSize(integrationRule.GetNPoints());
            data.quadratureWeights.SetSize(integrationRule.GetNPoints());
            data.cylindricalRadiusSquared.SetSize(integrationRule.GetNPoints());
        }
        int globalStellarElementCount = 0;
        MFEM_VERIFY(
            MPI_Allreduce(
                &localStellarElementCount, &globalStellarElementCount, 1, MPI_INT, MPI_SUM, m_fem.mesh->GetComm()
            ) == MPI_SUCCESS,
            "PreparedAngularMomentumOperator could not count stellar elements."
        );
        MFEM_VERIFY(globalStellarElementCount > 0, "PreparedAngularMomentumOperator found no stellar elements.");
    }

    std::optional<mapping::MappingStatus>
    PreparedAngularMomentumOperator::RefreshGeometry(const mfem::Vector &displacement) {
        MFEM_VERIFY(
            displacement.Size() == m_fem.displacementFes->GetTrueVSize(),
            "Angular-momentum geometry has the wrong displacement size."
        );
        if (!is_finite_vector(displacement)) {
            return mapping::MappingStatus::non_finite_input;
        }
        mfem::Vector displacementLocal;
        true_to_local(*m_fem.displacementFes, displacement, displacementLocal);
        if (!is_finite_vector(displacementLocal)) {
            return mapping::MappingStatus::non_finite_result;
        }
        mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());
        mapping::VolumeMappingContext mappingContext;

        for (ElementPAData &data : m_elements) {
            displacementLocal.GetSubVector(data.displacementDofs, data.baseDisplacement);
            m_fem.compactificationCoordinate->GetSubVector(data.compactificationDofs, data.compactification);
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(data.baseDisplacement);
            }
            if (data.compactificationDofTransformation != nullptr) {
                data.compactificationDofTransformation->InvTransformPrimal(data.compactification);
            }
            if (!is_finite_vector(data.baseDisplacement)) {
                return mapping::MappingStatus::non_finite_result;
            }
            MFEM_VERIFY(
                is_finite_vector(data.compactification),
                "Angular-momentum preparation encountered invalid static compactification data."
            );
            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(data.elementId);
            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);
            const mapping::ElementDisplacementData displacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, data.compactification
            );
            const mapping::ElementMappingData mappingData{
                .displacement = displacementData, .compactification = compactificationData
            };
            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);
            for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
                const mapping::MappingStatus status = m_domainMapper.EvaluateVolume(
                    mappingData, *transformation, data.integrationRule->IntPoint(quadraturePoint), workspace,
                    mappingContext
                );
                if (status != mapping::MappingStatus::valid) {
                    return status;
                }
                if (mappingContext.mapping.compactified) {
                    return mapping::MappingStatus::at_compactified_infinity;
                }
                data.cylindricalRadiusSquared(quadraturePoint) =
                    CylindricalRadiusSquared(mappingContext.mapping.physical_position);
                if (!std::isfinite(data.cylindricalRadiusSquared(quadraturePoint))) {
                    return mapping::MappingStatus::non_finite_result;
                }
                data.mappingContexts.Store(quadraturePoint, mappingContext);
                data.quadratureWeights(quadraturePoint) = mappingContext.quadrature.weight;
            }
        }
        return std::nullopt;
    }

    bool PreparedAngularMomentumOperator::RefreshDensity(const mfem::Vector &density) {
        MFEM_VERIFY(density.Size() == m_fem.densityFes->GetTrueVSize(), "Angular-momentum density has the wrong size.");
        if (!is_finite_vector(density)) {
            return false;
        }
        mfem::Vector densityLocal;
        true_to_local(*m_fem.densityFes, density, densityLocal);
        if (!is_finite_vector(densityLocal)) {
            return false;
        }
        mfem::Vector elementDensity;
        for (ElementPAData &data : m_elements) {
            densityLocal.GetSubVector(data.densityDofs, elementDensity);
            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(elementDensity);
            }
            if (!is_finite_vector(elementDensity)) {
                return false;
            }
            data.densityBasis->GetValues().Mult(elementDensity, data.density);
            if (!is_finite_vector(data.density)) {
                return false;
            }
        }
        return true;
    }

    std::optional<AngularMomentumPreparationRejection> PreparedAngularMomentumOperator::TryAssembleResidual() {
        double localMomentOfInertia = 0.0;
        for (const ElementPAData &data : m_elements) {
            for (int quadraturePoint = 0; quadraturePoint < data.density.Size(); ++quadraturePoint) {
                localMomentOfInertia += data.density(quadraturePoint) * data.cylindricalRadiusSquared(quadraturePoint) *
                                        data.quadratureWeights(quadraturePoint);
            }
        }
        m_momentOfInertia = GlobalSum(localMomentOfInertia);
        if (!std::isfinite(m_momentOfInertia)) {
            return AngularMomentumPreparationRejection{
                .reason          = AngularMomentumPreparationRejectionReason::non_finite_moment_of_inertia,
                .momentOfInertia = m_momentOfInertia
            };
        }
        if (m_momentOfInertia < 0.0) {
            return AngularMomentumPreparationRejection{
                .reason          = AngularMomentumPreparationRejectionReason::negative_moment_of_inertia,
                .momentOfInertia = m_momentOfInertia
            };
        }
        MFEM_VERIFY(
            std::isfinite(m_constraint.targetAngularMomentum().value()),
            "PreparedAngularMomentumOperator has a non-finite configured target angular momentum."
        );
        m_currentAngularMomentum = m_angularVelocity * m_momentOfInertia;
        m_cachedResidual.SetSize(1);
        m_cachedResidual(0) = m_currentAngularMomentum - m_constraint.targetAngularMomentum().value();
        if (!std::isfinite(m_currentAngularMomentum) || !std::isfinite(m_cachedResidual(0))) {
            return AngularMomentumPreparationRejection{
                .reason          = AngularMomentumPreparationRejectionReason::non_finite_residual,
                .momentOfInertia = m_momentOfInertia
            };
        }
        ++m_preparationCount;
        return std::nullopt;
    }

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

    double
    PreparedAngularMomentumOperator::EvaluateDensityMomentActionLocal(const mfem::Vector &densityVariation) const {
        MFEM_VERIFY(
            densityVariation.Size() == m_fem.densityFes->GetTrueVSize(),
            "Angular-momentum density action has the wrong true-vector size."
        );
        true_to_local(*m_fem.densityFes, densityVariation, m_densityVariationLocal);
        mfem::Vector quadratureDensityVariation;
        double localAction = 0.0;
        for (const ElementPAData &data : m_elements) {
            m_densityVariationLocal.GetSubVector(data.densityDofs, m_elementDensityVariation);
            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(m_elementDensityVariation);
            }
            quadratureDensityVariation.SetSize(data.integrationRule->GetNPoints());
            data.densityBasis->GetValues().Mult(m_elementDensityVariation, quadratureDensityVariation);
            for (int quadraturePoint = 0; quadraturePoint < quadratureDensityVariation.Size(); ++quadraturePoint) {
                localAction += quadratureDensityVariation(quadraturePoint) *
                               data.cylindricalRadiusSquared(quadraturePoint) * data.quadratureWeights(quadraturePoint);
            }
        }
        return localAction;
    }

    double PreparedAngularMomentumOperator::EvaluateDisplacementMomentActionLocal(
        const mfem::Vector &displacementVariation
    ) const {
        MFEM_VERIFY(
            displacementVariation.Size() == m_fem.displacementFes->GetTrueVSize(),
            "Angular-momentum displacement action has the wrong true-vector size."
        );
        true_to_local(*m_fem.displacementFes, displacementVariation, m_displacementVariationLocal);
        mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());
        mapping::VolumeMappingVariation variation;
        mapping::VolumeMappingContext mappingContext;
        double localAction = 0.0;
        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 mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);
            const mapping::ElementDisplacementData baseDisplacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
            const mapping::ElementDisplacementData directionData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, m_elementDisplacementVariation);
            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, data.compactification
            );
            const mapping::ElementMappingData mappingData{
                .displacement = baseDisplacementData, .compactification = compactificationData
            };
            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);
            for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
                data.mappingContexts.Load(quadraturePoint, mappingContext);
                const mapping::MappingStatus status = m_domainMapper.EvaluateVolumeVariation(
                    mappingData, directionData, *transformation, data.integrationRule->IntPoint(quadraturePoint),
                    mappingContext, workspace, variation
                );
                MFEM_VERIFY(
                    status == mapping::MappingStatus::valid,
                    "Mapped angular-momentum variation is invalid. Element: " << data.elementId
                );
                const double radiusSquaredVariation = CylindricalRadiusSquaredVariation(
                    mappingContext.mapping.physical_position, variation.mapping.physical_position_variation
                );
                localAction += data.density(quadraturePoint) *
                               (radiusSquaredVariation * data.quadratureWeights(quadraturePoint) +
                                data.cylindricalRadiusSquared(quadraturePoint) * variation.weight_variation);
            }
        }
        return localAction;
    }

    void PreparedAngularMomentumOperator::ApplyDensityJacobianAction(
        const mfem::Vector &densityVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(
            densityVariation.Size() == m_gravityContext.GetDensityMap().reduced_size(),
            "Angular-momentum density action has the wrong reduced size."
        );
        verify_finite_vector(densityVariation, "Angular-momentum density direction is non-finite.");
        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        action.SetSize(1);
        action(0) = m_angularVelocity * GlobalSum(EvaluateDensityMomentActionLocal(m_densityVariationTrue));
        ++m_actionStatistics.densityApplications;
    }

    void PreparedAngularMomentumOperator::ApplyDisplacementJacobianAction(
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(
            displacementVariation.Size() == m_gravityContext.GetDisplacementMap().reduced_size(),
            "Angular-momentum displacement action has the wrong reduced size."
        );
        verify_finite_vector(displacementVariation, "Angular-momentum displacement direction is non-finite.");
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
        action.SetSize(1);
        action(0) = m_angularVelocity * GlobalSum(EvaluateDisplacementMomentActionLocal(m_displacementVariationTrue));
        ++m_actionStatistics.displacementApplications;
    }

    void PreparedAngularMomentumOperator::ApplyAngularVelocityJacobianAction(
        const double angularVelocityVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(std::isfinite(angularVelocityVariation), "Angular-velocity direction is non-finite.");
        action.SetSize(1);
        action(0) = m_momentOfInertia * angularVelocityVariation;
        ++m_actionStatistics.angularVelocityApplications;
    }

    void PreparedAngularMomentumOperator::ApplyCompleteJacobianAction(
        const mfem::Vector &densityVariation,
        const mfem::Vector &displacementVariation,
        const double angularVelocityVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(
            densityVariation.Size() == m_gravityContext.GetDensityMap().reduced_size() &&
                displacementVariation.Size() == m_gravityContext.GetDisplacementMap().reduced_size(),
            "Angular-momentum complete action has incompatible reduced coordinates."
        );
        verify_finite_vector(densityVariation, "Angular-momentum density direction is non-finite.");
        verify_finite_vector(displacementVariation, "Angular-momentum displacement direction is non-finite.");
        MFEM_VERIFY(std::isfinite(angularVelocityVariation), "Angular-velocity direction is non-finite.");
        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
        const double localMomentAction = EvaluateDensityMomentActionLocal(m_densityVariationTrue) +
                                         EvaluateDisplacementMomentActionLocal(m_displacementVariationTrue);
        action.SetSize(1);
        action(0) = m_angularVelocity * GlobalSum(localMomentAction) + m_momentOfInertia * angularVelocityVariation;
        ++m_actionStatistics.completeApplications;
    }

    double
    PreparedAngularMomentumOperator::CylindricalRadiusSquared(const mfem::Vector &physicalPosition) const noexcept {
        const auto &axis     = m_constraint.specification().axis();
        const auto &center   = m_constraint.specification().center();
        double radiusSquared = 0.0;
        double axialPosition = 0.0;
        for (int component = 0; component < 3; ++component) {
            const double relative = physicalPosition(component) - center[static_cast<std::size_t>(component)];
            radiusSquared += relative * relative;
            axialPosition += axis[static_cast<std::size_t>(component)] * relative;
        }
        const double perpendicularRadiusSquared = radiusSquared - axialPosition * axialPosition;
        if (!std::isfinite(perpendicularRadiusSquared)) {
            return perpendicularRadiusSquared;
        }
        return std::max(0.0, perpendicularRadiusSquared);
    }

    double PreparedAngularMomentumOperator::CylindricalRadiusSquaredVariation(
        const mfem::Vector &physicalPosition,
        const mfem::Vector &physicalPositionVariation
    ) const noexcept {
        const auto &axis            = m_constraint.specification().axis();
        const auto &center          = m_constraint.specification().center();
        double relativeDotVariation = 0.0;
        double axialPosition        = 0.0;
        double axialVariation       = 0.0;
        for (int component = 0; component < 3; ++component) {
            const double relative = physicalPosition(component) - center[static_cast<std::size_t>(component)];
            relativeDotVariation += relative * physicalPositionVariation(component);
            axialPosition += axis[static_cast<std::size_t>(component)] * relative;
            axialVariation += axis[static_cast<std::size_t>(component)] * physicalPositionVariation(component);
        }
        return 2.0 * (relativeDotVariation - axialPosition * axialVariation);
    }

    double PreparedAngularMomentumOperator::GlobalSum(const double localValue) const {
        double globalValue = 0.0;
        if (MPI_Allreduce(&localValue, &globalValue, 1, MPI_DOUBLE, MPI_SUM, m_fem.mesh->GetComm()) != MPI_SUCCESS) {
            throw std::runtime_error("PreparedAngularMomentumOperator could not reduce the moment of inertia.");
        }
        return globalValue;
    }

    bool PreparedAngularMomentumOperator::IsPrepared() const noexcept {
        if (!m_isPrepared || !m_gravityContext.IsPrepared()) {
            return false;
        }
        const auto &revisions = m_gravityContext.GetRevisions();
        return revisions.discretization.value == m_preparedDependencies.discretization.revision &&
               revisions.density.value == m_preparedDependencies.density.revision &&
               revisions.displacement.value == m_preparedDependencies.displacement.revision;
    }

    double PreparedAngularMomentumOperator::GetMomentOfInertia() const {
        VerifyPrepared();
        return m_momentOfInertia;
    }

    double PreparedAngularMomentumOperator::GetAngularVelocity() const {
        VerifyPrepared();
        return m_angularVelocity;
    }

    double PreparedAngularMomentumOperator::GetCurrentAngularMomentum() const {
        VerifyPrepared();
        return m_currentAngularMomentum;
    }

    double PreparedAngularMomentumOperator::GetTargetAngularMomentum() const noexcept {
        return m_constraint.targetAngularMomentum().value();
    }

    physics::RigidRotation PreparedAngularMomentumOperator::GetRotation() const {
        VerifyPrepared();
        return m_constraint.makeRotation(m_angularVelocity);
    }

    AngularMomentumConstraintReport PreparedAngularMomentumOperator::GetConstraintReport() const {
        VerifyPrepared();
        const double target   = GetTargetAngularMomentum();
        const double residual = m_currentAngularMomentum - target;
        return {
            .targetAngularMomentum   = target,
            .achievedAngularMomentum = m_currentAngularMomentum,
            .momentOfInertia         = m_momentOfInertia,
            .angularVelocity         = m_angularVelocity,
            .dimensionalResidual     = residual,
            .scaledResidual          = residual / std::max(std::abs(target), 1.0e-300)
        };
    }

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

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

    const PreparedAngularMomentumActionStatistics &
    PreparedAngularMomentumOperator::GetActionStatistics() const noexcept {
        return m_actionStatistics;
    }

    const models::CompiledFixedAngularMomentum &
    PreparedAngularMomentumOperator::GetCompiledConstraint() const noexcept {
        return m_constraint;
    }

    void PreparedAngularMomentumOperator::VerifyPrepared() const {
        MFEM_VERIFY(IsPrepared(), "The angular-momentum invariant must be prepared before application.");
    }
} // namespace mean_field::operators