module;
#include "profile.h"
#include <array>
#include <cmath>
#include <cstdint>
#include <expected>
#include <memory>
#include <mfem.hpp>
#include <optional>
#include <stdexcept>
#include <string>

#include <mpi.h>

module mean_field;

import :fem.reference_tables;
import :operators.prepared_hdiv_mass;

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

    [[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]] mean_field::operators::HDivMassPreparationResult synchronize_preparation_failure(
        const mean_field::mapping::MappingStatus localMappingStatus,
        const bool localNonFiniteArithmetic,
        const MPI_Comm communicator
    ) {
        std::array<int, 3> localFailures{0, 0, localNonFiniteArithmetic ? 1 : 0};
        if (localMappingStatus != mean_field::mapping::MappingStatus::valid) {
            const int encodedStatus = static_cast<int>(localMappingStatus) + 1;
            localFailures[is_candidate_mapping_failure(localMappingStatus) ? 0 : 1] = encodedStatus;
        }

        std::array<int, 3> globalFailures{};
        if (MPI_Allreduce(
                localFailures.data(), globalFailures.data(), static_cast<int>(localFailures.size()), MPI_INT, MPI_MAX,
                communicator
            ) != MPI_SUCCESS) {
            throw std::runtime_error("PreparedMappedHDivMassOperator could not synchronize candidate validity.");
        }
        if (globalFailures[1] != 0) {
            throw std::runtime_error(
                "PreparedMappedHDivMassOperator encountered a structural mapping failure with status " +
                std::to_string(globalFailures[1] - 1) + "."
            );
        }
        if (globalFailures[0] != 0) {
            return std::unexpected(
                mean_field::operators::HDivMassPreparationRejection{
                    .reason        = mean_field::operators::HDivMassPreparationRejectionReason::invalid_mapping,
                    .mappingStatus = static_cast<mean_field::mapping::MappingStatus>(globalFailures[0] - 1)
                }
            );
        }
        if (globalFailures[2] != 0) {
            return std::unexpected(
                mean_field::operators::HDivMassPreparationRejection{
                    .reason = mean_field::operators::HDivMassPreparationRejectionReason::non_finite_arithmetic
                }
            );
        }
        return {};
    }

    [[nodiscard]] mean_field::mapping::MappingStatus higher_priority_mapping_status(
        const mean_field::mapping::MappingStatus left,
        const mean_field::mapping::MappingStatus right
    ) noexcept {
        if (left == mean_field::mapping::MappingStatus::valid) {
            return right;
        }
        if (right == mean_field::mapping::MappingStatus::valid) {
            return left;
        }
        const bool leftIsCandidate  = is_candidate_mapping_failure(left);
        const bool rightIsCandidate = is_candidate_mapping_failure(right);
        if (leftIsCandidate != rightIsCandidate) {
            return leftIsCandidate ? right : left;
        }
        return static_cast<int>(right) > static_cast<int>(left) ? right : left;
    }

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

    int get_operator_size(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "PreparedMappedHDivMassOperator requires the "
                                         "gravity-gradient finite-element space."
        );
        return mean_field::field::make_field_dof_map<mean_field::field::Gravity, DomainSchema>(*f.gravityFluxFes)
            .reduced_size();
    }

    bool communicator_has_single_rank(const MPI_Comm communicator) {
        int size = 0;
        MFEM_VERIFY(MPI_Comm_size(communicator, &size) == MPI_SUCCESS, "Failed to query the MPI communicator size.");
        MFEM_VERIFY(size > 0, "The MPI communicator must contain at least one rank.");
        return size == 1;
    }

    void true_to_local(
        const mfem::ParFiniteElementSpace &finite_element_space,
        const mfem::Vector &true_vector,
        mfem::Vector &local_vector
    ) {
        local_vector.SetSize(finite_element_space.GetVSize());

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

        if (prolongation != nullptr) {
            prolongation->Mult(true_vector, local_vector);
        } else {
            local_vector = true_vector;
        }
    }

    void local_to_true(
        const mfem::ParFiniteElementSpace &finite_element_space,
        const mfem::Vector &local_vector,
        mfem::Vector &true_vector
    ) {
        true_vector.SetSize(finite_element_space.GetTrueVSize());
        true_vector                        = 0.0;
        const mfem::Operator *prolongation = finite_element_space.GetProlongationMatrix();
        if (prolongation != nullptr) {
            prolongation->MultTranspose(local_vector, true_vector);
        } else {
            true_vector = local_vector;
        }
    }

    mean_field::quadrature::MappingKind get_mapping_kind(
        const mean_field::mapping::DomainMapper &domain_mapper,
        const mfem::ElementTransformation &transformation
    ) {
        return domain_mapper.IsCompactifiedElement(transformation) ? mean_field::quadrature::MappingKind::kelvin
                                                                   : mean_field::quadrature::MappingKind::general;
    }

    const mfem::IntegrationRule &get_hdiv_mass_rule(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapper &domain_mapper,
        const mfem::FiniteElement &element,
        const mfem::ElementTransformation &transformation
    ) {
        using GravityField = mean_field::field::Field<mean_field::field::Gravity>;
        const mean_field::quadrature::Query query =
            GravityField::make_query<mean_field::field::Gravity::Form::HDivMass>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
                mean_field::utils::DOMAINS::ALL, get_mapping_kind(domain_mapper, transformation)
            );
        const auto resolution = f.quadratureFactory->get(query, transformation.GetGeometryType());
        MFEM_VERIFY(
            resolution.integration_rule != nullptr,
            "The quadrature policy did not return an H(div) mass integration rule."
        );
        return *resolution.integration_rule;
    }

    int frozen_mapping_width(const int dimension) {
        return 3 * dimension + 4 * dimension * dimension + 3;
    }

    void freeze_mapping_context(
        const mean_field::mapping::VolumeMappingContext &context,
        const int quadrature_point,
        mfem::DenseMatrix &data
    ) {
        const int dimension                   = context.mapping.reference_position.Size();
        const int displacement_jacobian_start = 3 * dimension;
        const int mapping_jacobian_start      = displacement_jacobian_start + dimension * dimension;
        const int inverse_mapping_start       = mapping_jacobian_start + dimension * dimension;
        const int inverse_element_start       = inverse_mapping_start + dimension * dimension;
        const int scalar_start                = inverse_element_start + dimension * dimension;

        for (int component = 0; component < dimension; ++component) {
            data(quadrature_point, component)                 = context.mapping.reference_position(component);
            data(quadrature_point, dimension + component)     = context.mapping.displaced_position(component);
            data(quadrature_point, 2 * dimension + component) = context.mapping.physical_position(component);
        }
        for (int row = 0; row < dimension; ++row) {
            for (int column = 0; column < dimension; ++column) {
                const int entry = row * dimension + column;
                data(quadrature_point, displacement_jacobian_start + entry) =
                    context.mapping.displacement_jacobian(row, column);
                data(quadrature_point, mapping_jacobian_start + entry) = context.mapping.mapping_jacobian(row, column);
                data(quadrature_point, inverse_mapping_start + entry) =
                    context.mapping.inverse_mapping_jacobian(row, column);
                data(quadrature_point, inverse_element_start + entry) = context.quadrature.J_inv(row, column);
            }
        }
        data(quadrature_point, scalar_start)     = context.mapping.mapping_determinant;
        data(quadrature_point, scalar_start + 1) = context.quadrature.weight;
        data(quadrature_point, scalar_start + 2) = context.mapping.compactified ? 1.0 : 0.0;
    }

    void thaw_mapping_context(
        const mfem::DenseMatrix &data,
        const int quadrature_point,
        const int dimension,
        mean_field::mapping::VolumeMappingContext &context
    ) {
        const int displacement_jacobian_start = 3 * dimension;
        const int mapping_jacobian_start      = displacement_jacobian_start + dimension * dimension;
        const int inverse_mapping_start       = mapping_jacobian_start + dimension * dimension;
        const int inverse_element_start       = inverse_mapping_start + dimension * dimension;
        const int scalar_start                = inverse_element_start + dimension * dimension;

        context.mapping.reference_position.SetSize(dimension);
        context.mapping.displaced_position.SetSize(dimension);
        context.mapping.physical_position.SetSize(dimension);
        context.mapping.displacement_jacobian.SetSize(dimension, dimension);
        context.mapping.mapping_jacobian.SetSize(dimension, dimension);
        context.mapping.inverse_mapping_jacobian.SetSize(dimension, dimension);
        context.quadrature.J_inv.SetSize(dimension, dimension);

        for (int component = 0; component < dimension; ++component) {
            context.mapping.reference_position(component) = data(quadrature_point, component);
            context.mapping.displaced_position(component) = data(quadrature_point, dimension + component);
            context.mapping.physical_position(component)  = data(quadrature_point, 2 * dimension + component);
        }
        for (int row = 0; row < dimension; ++row) {
            for (int column = 0; column < dimension; ++column) {
                const int entry = row * dimension + column;
                context.mapping.displacement_jacobian(row, column) =
                    data(quadrature_point, displacement_jacobian_start + entry);
                context.mapping.mapping_jacobian(row, column) = data(quadrature_point, mapping_jacobian_start + entry);
                context.mapping.inverse_mapping_jacobian(row, column) =
                    data(quadrature_point, inverse_mapping_start + entry);
                context.quadrature.J_inv(row, column) = data(quadrature_point, inverse_element_start + entry);
            }
        }
        context.mapping.mapping_determinant = data(quadrature_point, scalar_start);
        context.mapping.compactified        = data(quadrature_point, scalar_start + 2) != 0.0;
        context.quadrature.detJ             = context.mapping.mapping_determinant;
        context.quadrature.weight           = data(quadrature_point, scalar_start + 1);
    }

    int find_representative_element(
        const mean_field::fem::FEM &f,
        const mfem::Array<int> &marker
    ) {
        for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
            const int attribute = f.mesh->GetAttribute(element_id);

            if (attribute > 0 && attribute <= marker.Size() && marker[attribute - 1] != 0) {
                return element_id;
            }
        }

        return -1;
    }

    void validate_uniform_domain_discretization(
        const mean_field::fem::FEM &f,
        const mfem::Array<int> &marker,
        const int representative_element_id
    ) {
        const mfem::FiniteElement &representative_element = *f.gravityFluxFes->GetFE(representative_element_id);
        const mfem::ElementTransformation &representative_transformation =
            *f.mesh->GetElementTransformation(representative_element_id);

        for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
            const int attribute = f.mesh->GetAttribute(element_id);

            if (attribute <= 0 || attribute > marker.Size() || marker[attribute - 1] == 0) {
                continue;
            }

            const mfem::FiniteElement &element                = *f.gravityFluxFes->GetFE(element_id);
            const mfem::ElementTransformation &transformation = *f.mesh->GetElementTransformation(element_id);

            MFEM_VERIFY(
                element.GetGeomType() == representative_element.GetGeomType(),
                "Prepared H(div) mass domains currently require a uniform "
                "element "
                "geometry."
            );
            MFEM_VERIFY(
                element.GetOrder() == representative_element.GetOrder(),
                "Prepared H(div) mass domains currently require a uniform "
                "finite-element order."
            );
            MFEM_VERIFY(
                transformation.OrderW() == representative_transformation.OrderW(),
                "Prepared H(div) mass domains currently require a uniform "
                "geometry-weight order."
            );
        }
    }

    class FrozenMappedHDivMassCoefficient final : public mfem::MatrixCoefficient {
    public:
        FrozenMappedHDivMassCoefficient(
            const mean_field::fem::FEM &f,
            const mean_field::mapping::DomainMapper &domain_mapper,
            const mfem::Vector &displacement_true,
            bool elevates_vacuum
        )
            : MatrixCoefficient(domain_mapper.GetDimension()),
              m_fem(f),
              m_domain_mapper(domain_mapper),
              m_workspace(domain_mapper.GetDimension()),
              m_elevates_vacuum(elevates_vacuum) {
            true_to_local(*m_fem.displacementFes, displacement_true, m_displacement_local);
        }

        void Eval(
            mfem::DenseMatrix &mass_tensor,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point
        ) override {
            transformation.SetIntPoint(&integration_point);

            const int element_id = transformation.ElementNo;
            MFEM_VERIFY(
                element_id >= 0 && element_id < m_fem.mesh->GetNE(),
                "Mapped H(div) mass coefficient received an invalid element ID."
            );

            const bool element_is_vacuum =
                DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Vacuum>(
                    transformation.Attribute
                );

            if (element_is_vacuum != m_elevates_vacuum) {
                mass_tensor.SetSize(m_domain_mapper.GetDimension());
                mass_tensor = 0.0;
                return;
            }

            LoadElement(element_id);

            const mean_field::mapping::ElementMappingData mapping_data{
                .displacement = *m_displacement_data, .compactification = *m_compactification_data
            };

            mean_field::mapping::VolumeMappingContext mapping_context;

            const mean_field::mapping::MappingStatus status = m_domain_mapper.EvaluateVolume(
                mapping_data, transformation, integration_point, m_workspace, mapping_context
            );

            if (status != mean_field::mapping::MappingStatus::valid) {
                m_mappingFailure = higher_priority_mapping_status(m_mappingFailure, status);
                mass_tensor.SetSize(m_domain_mapper.GetDimension());
                mass_tensor = 0.0;
                return;
            }

            const mfem::DenseMatrix &mapping_jacobian = mapping_context.mapping.mapping_jacobian;
            const double mapping_determinant          = mapping_context.mapping.mapping_determinant;

            mfem::MultAtB(mapping_jacobian, mapping_jacobian, mass_tensor);
            mass_tensor *= 1.0 / mapping_determinant;
            if (!matrix_is_finite(mass_tensor)) {
                m_nonFiniteArithmetic = true;
                mass_tensor           = 0.0;
            }
        }

        [[nodiscard]] mean_field::mapping::MappingStatus GetMappingFailure() const noexcept {
            return m_mappingFailure;
        }

        [[nodiscard]] bool HasNonFiniteArithmetic() const noexcept {
            return m_nonFiniteArithmetic;
        }

    private:
        void LoadElement(const int element_id) {
            if (element_id == m_cached_element_id) {
                return;
            }

            const mfem::FiniteElement &displacement_element     = *m_fem.displacementFes->GetFE(element_id);
            const mfem::FiniteElement &compactification_element = *m_fem.compactificationFes->GetFE(element_id);

            mfem::DofTransformation *displacement_dof_transformation =
                m_fem.displacementFes->GetElementVDofs(element_id, m_displacement_dofs);
            mfem::DofTransformation *compactification_dof_transformation =
                m_fem.compactificationFes->GetElementDofs(element_id, m_compactification_dofs);

            m_displacement_local.GetSubVector(m_displacement_dofs, m_element_displacement);
            m_fem.compactificationCoordinate->GetSubVector(m_compactification_dofs, m_element_compactification);

            if (displacement_dof_transformation != nullptr) {
                displacement_dof_transformation->InvTransformPrimal(m_element_displacement);
            }

            if (compactification_dof_transformation != nullptr) {
                compactification_dof_transformation->InvTransformPrimal(m_element_compactification);
            }

            m_displacement_data = std::make_unique<mean_field::mapping::ElementDisplacementData>(
                mean_field::mapping::ElementDisplacementDataFromElementVDofs(
                    displacement_element, m_element_displacement
                )
            );

            m_compactification_data = std::make_unique<mean_field::mapping::ElementCompactificationData>(
                compactification_element, m_element_compactification
            );

            m_cached_element_id = element_id;
        }

        const mean_field::fem::FEM &m_fem;
        const mean_field::mapping::DomainMapper &m_domain_mapper;

        mfem::Vector m_displacement_local;

        mfem::Array<int> m_displacement_dofs;
        mfem::Array<int> m_compactification_dofs;

        mfem::Vector m_element_displacement;
        mfem::Vector m_element_compactification;

        std::unique_ptr<mean_field::mapping::ElementDisplacementData> m_displacement_data;
        std::unique_ptr<mean_field::mapping::ElementCompactificationData> m_compactification_data;

        mean_field::mapping::DomainMapper::Workspace m_workspace;
        int m_cached_element_id{-1};
        bool m_elevates_vacuum;
        mean_field::mapping::MappingStatus m_mappingFailure{mean_field::mapping::MappingStatus::valid};
        bool m_nonFiniteArithmetic{false};
    };
} // namespace

namespace mean_field::operators {
    PreparedMappedHDivMassOperator::PreparedMappedHDivMassOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper
    )
        : Operator(get_operator_size(f)),
          m_fem(f),
          m_domain_mapper(domain_mapper),
          m_flux_map(
              field::make_field_dof_map<
                  field::Gravity,
                  DomainSchema>(*f.gravityFluxFes)
          ),
          m_displacement_map(
              field::make_field_dof_map<
                  field::Displacement,
                  DomainSchema>(*f.displacementFes)
          ),
          m_variationWorkspace(domain_mapper.GetDimension()),
          m_single_rank(communicator_has_single_rank(f.gravityFluxFes->GetComm())) {
        MFEM_VERIFY(f.mesh != nullptr, "PreparedMappedHDivMassOperator requires a mesh.");
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "PreparedMappedHDivMassOperator requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "PreparedMappedHDivMassOperator requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "PreparedMappedHDivMassOperator requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "PreparedMappedHDivMassOperator requires the compactification "
                                                     "coordinate."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "PreparedMappedHDivMassOperator requires the quadrature-rule "
                                            "factory."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "The stateless domain-mapper dimension does not match the mesh "
            "dimension."
        );

        m_stellar_marker = utils::domain::make_attribute_marker<utils::domain::Stellar, DomainSchema>(*f.mesh);
        m_vacuum_marker  = utils::domain::make_attribute_marker<utils::domain::Vacuum, DomainSchema>(*f.mesh);

        const int stellar_element_id = find_representative_element(f, m_stellar_marker);
        const int vacuum_element_id  = find_representative_element(f, m_vacuum_marker);

        MFEM_VERIFY(
            stellar_element_id >= 0, "PreparedMappedHDivMassOperator requires "
                                     "at least one stellar element."
        );
        MFEM_VERIFY(
            vacuum_element_id >= 0, "PreparedMappedHDivMassOperator requires at "
                                    "least one compactified vacuum element."
        );

        validate_uniform_domain_discretization(f, m_stellar_marker, stellar_element_id);
        validate_uniform_domain_discretization(f, m_vacuum_marker, vacuum_element_id);
    }

    mapping::MappingStatus PreparedMappedHDivMassOperator::PrepareVariationData() {
        MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::PrepareVariationData", 0);

        m_variationElements.clear();
        m_variationElements.reserve(m_fem.mesh->GetNE());

        mfem::Vector displacementLocal;
        true_to_local(*m_fem.displacementFes, m_displacement_true, displacementLocal);

        mfem::Vector elementDisplacement;
        mfem::Vector elementCompactification;
        mapping::VolumeMappingContext mappingContext;

        for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
            m_variationElements.emplace_back();
            ElementVariationData &data = m_variationElements.back();
            data.elementId             = elementId;

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

            displacementLocal.GetSubVector(data.displacementDofs, elementDisplacement);
            m_fem.compactificationCoordinate->GetSubVector(data.compactificationDofs, elementCompactification);
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(elementDisplacement);
            }
            if (compactificationDofTransformation != nullptr) {
                compactificationDofTransformation->InvTransformPrimal(elementCompactification);
            }
            data.baseDisplacement                              = elementDisplacement;
            data.compactification                              = elementCompactification;

            const mfem::FiniteElement &gravityGradientElement  = *m_fem.gravityFluxFes->GetFE(elementId);
            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(elementId);
            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(elementId);
            mfem::ElementTransformation *transformation        = m_fem.mesh->GetElementTransformation(elementId);
            MFEM_VERIFY(
                transformation != nullptr, "Prepared H(div) variation data received a null element transformation."
            );

            data.integrationRule = &get_hdiv_mass_rule(m_fem, m_domain_mapper, gravityGradientElement, *transformation);
            const int dimension  = m_domain_mapper.GetDimension();
            if (gravityGradientElement.GetMapType() == mfem::FiniteElement::H_DIV &&
                gravityGradientElement.GetDim() == dimension && gravityGradientElement.GetRangeDim() == dimension &&
                transformation->GetSpaceDim() == dimension) {
                data.gravityReferenceTable =
                    m_fem.GetReferenceTables().GetVectorTable(gravityGradientElement, *data.integrationRule);
                data.meshPiolaJacobians.SetSize(data.integrationRule->GetNPoints(), dimension * dimension);
                data.referenceWeights.SetSize(data.integrationRule->GetNPoints());
            }
            data.frozenMappingData.SetSize(
                data.integrationRule->GetNPoints(), frozen_mapping_width(m_domain_mapper.GetDimension())
            );

            const mapping::ElementDisplacementData displacementData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, data.compactification
            );
            const mapping::ElementMappingData mappingData{
                .displacement = displacementData, .compactification = compactificationData
            };

            for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
                const mapping::MappingStatus status            = m_domain_mapper.EvaluateVolume(
                    mappingData, *transformation, integrationPoint, m_variationWorkspace, mappingContext
                );
                if (status != mapping::MappingStatus::valid) {
                    return status;
                }
                freeze_mapping_context(mappingContext, quadraturePoint, data.frozenMappingData);
                if (data.gravityReferenceTable != nullptr) {
                    // CalcVShape_RT = reference_shape * J_mesh^T / Weight.
                    // Cache only this small factor, never the mapped basis.
                    const double meshWeight                = transformation->Weight();
                    const mfem::DenseMatrix &meshJacobian  = transformation->Jacobian();
                    const double inverseMeshWeight         = 1.0 / meshWeight;
                    data.referenceWeights(quadraturePoint) = integrationPoint.weight * meshWeight;
                    for (int row = 0; row < dimension; ++row) {
                        for (int column = 0; column < dimension; ++column) {
                            const double entry = inverseMeshWeight * meshJacobian(row, column);
                            if (!std::isfinite(entry))
                                return mapping::MappingStatus::non_finite_result;
                            data.meshPiolaJacobians(quadraturePoint, row * dimension + column) = entry;
                        }
                    }
                    if (!std::isfinite(data.referenceWeights(quadraturePoint))) {
                        return mapping::MappingStatus::non_finite_result;
                    }
                }
            }
        }
        return mapping::MappingStatus::valid;
    }

    void PreparedMappedHDivMassOperator::Prepare(const mfem::Vector &displacement) {
        MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::Prepare linearization", 0);
        auto result = TryPrepareImpl(displacement, PreparationMode::linearization);
        if (!result.has_value()) {
            throwHDivMassPreparationRejection(result.error());
        }
    }

    void PreparedMappedHDivMassOperator::PreparePrimal(const mfem::Vector &displacement) {
        MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::Prepare primal", 0);
        auto result = TryPrepareImpl(displacement, PreparationMode::primal);
        if (!result.has_value()) {
            throwHDivMassPreparationRejection(result.error());
        }
    }

    HDivMassPreparationResult PreparedMappedHDivMassOperator::TryPrepare(const mfem::Vector &displacement) {
        MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::TryPrepare linearization", 0);
        return TryPrepareImpl(displacement, PreparationMode::linearization);
    }

    HDivMassPreparationResult PreparedMappedHDivMassOperator::TryPreparePrimal(const mfem::Vector &displacement) {
        MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::TryPrepare primal", 0);
        return TryPrepareImpl(displacement, PreparationMode::primal);
    }

    HDivMassPreparationResult PreparedMappedHDivMassOperator::TryPrepareImpl(
        const mfem::Vector &displacement,
        const PreparationMode mode
    ) {
        MFEM_VERIFY(
            displacement.Size() == m_displacement_map.reduced_size(),
            "PreparedMappedHDivMassOperator received a displacement vector "
            "with "
            "the wrong size."
        );

        bool localNonFiniteInput = false;
        for (int i = 0; i < displacement.Size(); ++i) {
            localNonFiniteInput = localNonFiniteInput || !std::isfinite(displacement(i));
        }
        if (auto inputResult = synchronize_preparation_failure(
                localNonFiniteInput ? mapping::MappingStatus::non_finite_input : mapping::MappingStatus::valid, false,
                m_fem.mesh->GetComm()
            );
            !inputResult.has_value()) {
            m_is_prepared        = false;
            m_has_variation_data = false;
            return inputResult;
        }

        m_is_prepared        = false;
        m_has_variation_data = false;

        m_displacement_true.SetSize(m_displacement_map.full_size());
        m_displacement_map.scatter(displacement, m_displacement_true);

        const int stellar_element_id                        = find_representative_element(m_fem, m_stellar_marker);
        const int vacuum_element_id                         = find_representative_element(m_fem, m_vacuum_marker);

        const mfem::FiniteElement &stellar_element          = *m_fem.gravityFluxFes->GetFE(stellar_element_id);
        const mfem::FiniteElement &vacuum_element           = *m_fem.gravityFluxFes->GetFE(vacuum_element_id);

        mfem::ElementTransformation &stellar_transformation = *m_fem.mesh->GetElementTransformation(stellar_element_id);
        mfem::ElementTransformation &vacuum_transformation  = *m_fem.mesh->GetElementTransformation(vacuum_element_id);

        m_stellar_mass_form.reset();
        m_vacuum_mass_form.reset();
        m_stellar_mass_coefficient.reset();
        m_vacuum_mass_coefficient.reset();

        auto stellarMassCoefficient =
            std::make_unique<FrozenMappedHDivMassCoefficient>(m_fem, m_domain_mapper, m_displacement_true, false);
        auto vacuumMassCoefficient =
            std::make_unique<FrozenMappedHDivMassCoefficient>(m_fem, m_domain_mapper, m_displacement_true, true);
        auto *stellarMassCoefficientView = stellarMassCoefficient.get();
        auto *vacuumMassCoefficientView  = vacuumMassCoefficient.get();
        m_stellar_mass_coefficient       = std::move(stellarMassCoefficient);
        m_vacuum_mass_coefficient        = std::move(vacuumMassCoefficient);

        m_stellar_mass_form              = std::make_unique<mfem::ParBilinearForm>(m_fem.gravityFluxFes.get());
        m_vacuum_mass_form               = std::make_unique<mfem::ParBilinearForm>(m_fem.gravityFluxFes.get());
        m_stellar_mass_form->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);
        m_vacuum_mass_form->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);

        auto stellar_integrator = std::make_unique<mfem::VectorFEMassIntegrator>(*m_stellar_mass_coefficient);
        auto vacuum_integrator  = std::make_unique<mfem::VectorFEMassIntegrator>(*m_vacuum_mass_coefficient);

        m_fem.quadratureFactory->configure_gravity_hdiv_mass(
            *stellar_integrator, quadrature::QuadratureRole::discretization, stellar_element, stellar_transformation,
            utils::DOMAINS::STELLAR, quadrature::MappingKind::general
        );

        m_fem.quadratureFactory->configure_gravity_hdiv_mass(
            *vacuum_integrator, quadrature::QuadratureRole::discretization, vacuum_element, vacuum_transformation,
            utils::DOMAINS::VACUUM, quadrature::MappingKind::kelvin
        );

        m_stellar_mass_form->AddDomainIntegrator(stellar_integrator.release(), m_stellar_marker);
        m_vacuum_mass_form->AddDomainIntegrator(vacuum_integrator.release(), m_vacuum_marker);
        m_stellar_mass_form->Assemble();
        m_vacuum_mass_form->Assemble();

        mapping::MappingStatus localMappingFailure = higher_priority_mapping_status(
            stellarMassCoefficientView->GetMappingFailure(), vacuumMassCoefficientView->GetMappingFailure()
        );
        bool localNonFiniteArithmetic =
            stellarMassCoefficientView->HasNonFiniteArithmetic() || vacuumMassCoefficientView->HasNonFiniteArithmetic();

        if (mode == PreparationMode::linearization) {
            if (localMappingFailure == mapping::MappingStatus::valid && !localNonFiniteArithmetic) {
                localMappingFailure = PrepareVariationData();
            }
        } else {
            m_variationElements.clear();
        }

        auto preparationResult =
            synchronize_preparation_failure(localMappingFailure, localNonFiniteArithmetic, m_fem.mesh->GetComm());
        if (!preparationResult.has_value()) {
            return preparationResult;
        }

        m_has_variation_data = mode == PreparationMode::linearization;
        m_is_prepared        = true;
        ++m_preparation_count;
        return {};
    }

    void PreparedMappedHDivMassOperator::Mult(
        const mfem::Vector &gravity_gradient,
        mfem::Vector &action
    ) const {
        MEAN_FIELD_PROFILE_SCOPE("PreparedMappedHDivMassOperator::Mult");

        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedHDivMassOperator must be prepared "
                           "before Mult is called."
        );
        MFEM_VERIFY(
            m_stellar_mass_form != nullptr && m_vacuum_mass_form != nullptr,
            "PreparedMappedHDivMassOperator has incomplete domain mass forms."
        );
        MFEM_VERIFY(
            gravity_gradient.Size() == Width(), "PreparedMappedHDivMassOperator received a gravity-gradient vector "
                                                "with the wrong size."
        );

        const mfem::Vector *gravity_gradient_true = &gravity_gradient;
        if (!m_flux_map.is_identity()) [[unlikely]] {
            m_flux_true.SetSize(m_flux_map.full_size());
            m_flux_map.scatter(gravity_gradient, m_flux_true);
            gravity_gradient_true = &m_flux_true;
        }

        mfem::Vector *action_true = &action;
        if (!m_flux_map.is_identity()) [[unlikely]] {
            m_action_true.SetSize(m_flux_map.full_size());
            action_true = &m_action_true;
        }

        if (m_single_rank) [[likely]] {
            action_true->SetSize(m_flux_map.full_size());
            m_domain_action_true.SetSize(m_flux_map.full_size());
            m_stellar_mass_form->Mult(*gravity_gradient_true, *action_true);
            m_vacuum_mass_form->Mult(*gravity_gradient_true, m_domain_action_true);
            *action_true += m_domain_action_true;
        } else {
            true_to_local(*m_fem.gravityFluxFes, *gravity_gradient_true, m_flux_local);
            m_action_local.SetSize(m_fem.gravityFluxFes->GetVSize());
            m_domain_action_local.SetSize(m_fem.gravityFluxFes->GetVSize());
            m_stellar_mass_form->Mult(m_flux_local, m_action_local);
            m_vacuum_mass_form->Mult(m_flux_local, m_domain_action_local);
            m_action_local += m_domain_action_local;
            local_to_true(*m_fem.gravityFluxFes, m_action_local, *action_true);
        }

        if (!m_flux_map.is_identity()) [[unlikely]] {
            action.SetSize(Height());
            m_flux_map.gather(m_action_true, action);
        }
    }

    void PreparedMappedHDivMassOperator::MultDisplacementVariationTrue(
        const mfem::Vector &gravityGradientTrue,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &actionVariationTrue
    ) const {
        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedHDivMassOperator must be prepared before applying a displacement variation."
        );
        MFEM_VERIFY(
            m_has_variation_data,
            "PreparedMappedHDivMassOperator requires linearization preparation before applying a displacement "
            "variation."
        );
        MFEM_VERIFY(
            gravityGradientTrue.Size() == m_fem.gravityFluxFes->GetTrueVSize(),
            "The full gravity-gradient vector has the wrong size."
        );
        MFEM_VERIFY(
            displacementVariationTrue.Size() == m_fem.displacementFes->GetTrueVSize(),
            "The full displacement variation has the wrong size."
        );

        true_to_local(*m_fem.gravityFluxFes, gravityGradientTrue, m_gravityGradientLocal);
        true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacementVariationLocal);
        m_localVariationAction.SetSize(m_fem.gravityFluxFes->GetVSize());
        m_localVariationAction = 0.0;

        const int dimension    = m_domain_mapper.GetDimension();

        for (const ElementVariationData &data : m_variationElements) {
            MFEM_VERIFY(
                data.integrationRule != nullptr &&
                    data.frozenMappingData.Height() == data.integrationRule->GetNPoints() &&
                    data.frozenMappingData.Width() == frozen_mapping_width(dimension),
                "Prepared H(div) variation data is incomplete."
            );

            m_gravityGradientLocal.GetSubVector(data.gravityGradientDofs, m_elementGravityGradient);
            m_displacementVariationLocal.GetSubVector(data.displacementDofs, m_elementDisplacementVariation);
            if (data.gravityGradientDofTransformation != nullptr) {
                data.gravityGradientDofTransformation->InvTransformPrimal(m_elementGravityGradient);
            }
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(m_elementDisplacementVariation);
            }

            const mfem::FiniteElement &gravityGradientElement  = *m_fem.gravityFluxFes->GetFE(data.elementId);
            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(data.elementId);
            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);
            mfem::ElementTransformation *transformation        = m_fem.mesh->GetElementTransformation(data.elementId);
            MFEM_VERIFY(
                transformation != nullptr,
                "Prepared H(div) displacement variation received a null element transformation."
            );

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

            m_elementVariationAction.SetSize(gravityGradientElement.GetDof());
            m_elementVariationAction = 0.0;
            m_gravityGradientValue.SetSize(dimension);
            m_gravityReferenceCellValue.SetSize(dimension);
            m_referenceCellDual.SetSize(dimension);
            m_massTensorVariationAction.SetSize(dimension);
            m_meshPiolaJacobian.SetSize(dimension, dimension);
            m_gravityGradientShape.SetSize(gravityGradientElement.GetDof(), dimension);
            m_massTensorVariation.SetSize(dimension, dimension);

            for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
                thaw_mapping_context(data.frozenMappingData, quadraturePoint, dimension, m_baseMappingContext);

                const mapping::MappingStatus status = m_domain_mapper.EvaluateVolumeVariation(
                    mappingData, directionData, *transformation, integrationPoint, m_baseMappingContext,
                    m_variationWorkspace, m_mappingVariation
                );
                MFEM_VERIFY(
                    status == mapping::MappingStatus::valid,
                    "Prepared H(div) displacement variation encountered an invalid mapping variation. Element: "
                        << data.elementId << ", quadrature point: " << quadraturePoint
                        << ", status: " << static_cast<int>(status)
                );

                mapping::ComputeHDivMassTensorVariation(
                    m_baseMappingContext.mapping, m_mappingVariation.mapping, m_massTensorVariation
                );

                if (data.gravityReferenceTable != nullptr) {
                    const mfem::DenseMatrix &referenceShape = data.gravityReferenceTable->GetValues(quadraturePoint);
                    referenceShape.MultTranspose(m_elementGravityGradient, m_gravityReferenceCellValue);
                    for (int row = 0; row < dimension; ++row) {
                        for (int column = 0; column < dimension; ++column) {
                            m_meshPiolaJacobian(row, column) =
                                data.meshPiolaJacobians(quadraturePoint, row * dimension + column);
                        }
                    }
                    m_meshPiolaJacobian.Mult(m_gravityReferenceCellValue, m_gravityGradientValue);
                    m_massTensorVariation.Mult(m_gravityGradientValue, m_massTensorVariationAction);
                    // Move the test-side Piola transform onto the three-vector
                    // dual before applying the reference basis transpose.
                    m_meshPiolaJacobian.MultTranspose(m_massTensorVariationAction, m_referenceCellDual);
                    referenceShape.AddMult(
                        m_referenceCellDual, m_elementVariationAction, data.referenceWeights(quadraturePoint)
                    );
                } else {
                    transformation->SetIntPoint(&integrationPoint);
                    gravityGradientElement.CalcVShape(*transformation, m_gravityGradientShape);
                    m_gravityGradientShape.MultTranspose(m_elementGravityGradient, m_gravityGradientValue);
                    m_massTensorVariation.Mult(m_gravityGradientValue, m_massTensorVariationAction);
                    const double referenceWeight = integrationPoint.weight * transformation->Weight();
                    m_gravityGradientShape.AddMult(
                        m_massTensorVariationAction, m_elementVariationAction, referenceWeight
                    );
                }
            }

            if (data.gravityGradientDofTransformation != nullptr) {
                data.gravityGradientDofTransformation->TransformDual(m_elementVariationAction);
            }
            m_localVariationAction.AddElementVector(data.gravityGradientDofs, m_elementVariationAction);
        }

        local_to_true(*m_fem.gravityFluxFes, m_localVariationAction, actionVariationTrue);
    }

    void PreparedMappedHDivMassOperator::AssembleDiagonal(mfem::Vector &diagonal) const {
        mfem::Vector true_diagonal;
        AssembleTrueDiagonal(true_diagonal);
        diagonal.SetSize(Height());
        m_flux_map.gather(true_diagonal, diagonal);
    }

    void PreparedMappedHDivMassOperator::AssembleTrueDiagonal(mfem::Vector &diagonal) const {
        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedHDivMassOperator must be prepared "
                           "before assembling its diagonal."
        );
        MFEM_VERIFY(
            m_stellar_mass_form != nullptr && m_vacuum_mass_form != nullptr,
            "PreparedMappedHDivMassOperator has incomplete domain mass forms."
        );

        diagonal.SetSize(m_flux_map.full_size());
        mfem::Vector domain_diagonal(m_flux_map.full_size());
        m_stellar_mass_form->AssembleDiagonal(diagonal);
        m_vacuum_mass_form->AssembleDiagonal(domain_diagonal);
        diagonal += domain_diagonal;
    }

    bool PreparedMappedHDivMassOperator::IsPrepared() const noexcept {
        return m_is_prepared;
    }

    bool PreparedMappedHDivMassOperator::HasVariationData() const noexcept {
        return m_has_variation_data;
    }

    std::uint64_t PreparedMappedHDivMassOperator::GetPreparationCount() const noexcept {
        return m_preparation_count;
    }

    const field::FieldDofMap &PreparedMappedHDivMassOperator::GetFluxMap() const noexcept {
        return m_flux_map;
    }

    const field::FieldDofMap &PreparedMappedHDivMassOperator::GetDisplacementMap() const noexcept {
        return m_displacement_map;
    }
} // namespace mean_field::operators