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

#include <mpi.h>

module mean_field;
import :operators.prepared_gravity_source;

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::GravitySourcePreparationResult 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("PreparedMappedGravitySourceOperator could not synchronize candidate validity.");
        }
        if (globalFailures[1] != 0) {
            throw std::runtime_error(
                "PreparedMappedGravitySourceOperator encountered a structural mapping failure with status " +
                std::to_string(globalFailures[1] - 1) + "."
            );
        }
        if (globalFailures[0] != 0) {
            return std::unexpected(
                mean_field::operators::GravitySourcePreparationRejection{
                    .reason        = mean_field::operators::GravitySourcePreparationRejectionReason::invalid_mapping,
                    .mappingStatus = static_cast<mean_field::mapping::MappingStatus>(globalFailures[0] - 1)
                }
            );
        }
        if (globalFailures[2] != 0) {
            return std::unexpected(
                mean_field::operators::GravitySourcePreparationRejection{
                    .reason = mean_field::operators::GravitySourcePreparationRejectionReason::non_finite_arithmetic
                }
            );
        }
        return {};
    }

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

    int get_operator_width(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.densityFes != nullptr, "PreparedMappedGravitySourceOperator requires the density "
                                     "finite-element space."
        );
        return mean_field::field::make_field_dof_map<mean_field::field::Density, DomainSchema>(*f.densityFes)
            .reduced_size();
    }

    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
    ) {
        MFEM_VERIFY(local_vector.Size() == finite_element_space.GetVSize(), "Local vector has the wrong size.");

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

    const mfem::IntegrationRule &get_source_rule(
        const mean_field::fem::FEM &f,
        const mfem::FiniteElement &density_element,
        const mfem::FiniteElement &potential_element,
        const mfem::ElementTransformation &transformation
    ) {
        using GravityField = mean_field::field::Field<mean_field::field::Gravity>;
        MFEM_VERIFY(
            density_element.GetOrder() == mean_field::field::Density::Scalar::familyOrder,
            "The prepared source trial element does not match the registered "
            "density field."
        );
        MFEM_VERIFY(
            potential_element.GetOrder() == mean_field::field::Gravity::Potential::familyOrder,
            "The prepared source test element does not match the registered "
            "gravity potential."
        );
        const mean_field::quadrature::Query query =
            GravityField::make_query<mean_field::field::Gravity::Form::SourceProjection>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
                mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
            );

        return *f.quadratureFactory->get(query, transformation.GetGeometryType()).integration_rule;
    }

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

        double Eval(
            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 gravity source coefficient received an invalid element "
                "ID."
            );
            if (DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Vacuum>(
                    transformation.Attribute
                )) {
                return 0.0;
            }

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

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

            if (status != mean_field::mapping::MappingStatus::valid) {
                m_mappingFailure = status;
                return 0.0;
            }
            const double mapping_determinant = m_mapping_context.mapping.mapping_determinant;

            m_inverse_element_jacobian       = m_mapping_context.quadrature.J_inv;

            const double value               = 4.0 * std::numbers::pi * mean_field::utils::G * mapping_determinant;
            if (!std::isfinite(value)) {
                m_nonFiniteArithmetic = true;
                return 0.0;
            }
            return value;
        }

        [[nodiscard]] const mfem::DenseMatrix &GetInverseElementJacobian() const noexcept {
            return m_inverse_element_jacobian;
        }

        [[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;
        mean_field::mapping::VolumeMappingContext m_mapping_context;
        mfem::DenseMatrix m_inverse_element_jacobian;
        int m_cached_element_id{-1};
        mean_field::mapping::MappingStatus m_mappingFailure{mean_field::mapping::MappingStatus::valid};
        bool m_nonFiniteArithmetic{false};
    };
} // namespace

namespace mean_field::operators {
    PreparedMappedGravitySourceOperator::PreparedMappedGravitySourceOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper
    )
        : Operator(
              get_operator_height(f),
              get_operator_width(f)
          ),
          m_fem(f),
          m_domain_mapper(domain_mapper),
          m_density_map(
              field::make_field_dof_map<
                  field::Density,
                  DomainSchema>(*f.densityFes)
          ),
          m_potential_map(
              field::make_field_dof_map<
                  field::Gravity,
                  DomainSchema>(*f.gravityPotentialFes)
          ),
          m_displacement_map(
              field::make_field_dof_map<
                  field::Displacement,
                  DomainSchema>(*f.displacementFes)
          ) {
        MFEM_VERIFY(f.mesh != nullptr, "PreparedMappedGravitySourceOperator requires a mesh.");
        MFEM_VERIFY(
            f.densityFes != nullptr, "PreparedMappedGravitySourceOperator requires the density "
                                     "finite-element space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "PreparedMappedGravitySourceOperator requires the "
                                              "gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "PreparedMappedGravitySourceOperator requires "
                                          "the displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "PreparedMappedGravitySourceOperator requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr,
            "PreparedMappedGravitySourceOperator requires the compactification "
            "coordinate."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "PreparedMappedGravitySourceOperator "
                                            "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);
    }

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

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

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

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

    GravitySourcePreparationResult PreparedMappedGravitySourceOperator::TryPrepareImpl(
        const mfem::Vector &displacement,
        const PreparationMode mode
    ) {
        MFEM_VERIFY(
            displacement.Size() == m_displacement_map.reduced_size(),
            "PreparedMappedGravitySourceOperator 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);
        m_elements.reserve(m_fem.mesh->GetNE());
        std::size_t prepared_element_count{0};

        FrozenMappedGravitySourceCoefficient source_coefficient(m_fem, m_domain_mapper, m_displacement_true);
        bool localNonFiniteQuadrature = false;

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

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

            if (prepared_element_count == m_elements.size()) {
                m_elements.emplace_back();
            }
            ElementPAData &data             = m_elements[prepared_element_count++];

            data.element_id                 = element_id;

            data.density_dof_transformation = m_fem.densityFes->GetElementDofs(element_id, data.density_dofs);

            data.potential_dof_transformation =
                m_fem.gravityPotentialFes->GetElementDofs(element_id, data.potential_dofs);

            if (mode == PreparationMode::linearization) {
                data.displacement_dof_transformation =
                    m_fem.displacementFes->GetElementVDofs(element_id, data.displacement_dofs);
            }

            const mfem::FiniteElement &density_element   = *m_fem.densityFes->GetFE(element_id);

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

            mfem::ElementTransformation &transformation  = *m_fem.mesh->GetElementTransformation(element_id);

            const mfem::IntegrationRule &integration_rule =
                get_source_rule(m_fem, density_element, potential_element, transformation);
            data.integration_rule            = &integration_rule;

            const int quadrature_point_count = integration_rule.GetNPoints();

            const int density_dof_count      = density_element.GetDof();

            const int potential_dof_count    = potential_element.GetDof();

            if (density_element.GetMapType() == mfem::FiniteElement::VALUE) {
                data.density_reference = m_fem.GetReferenceTables().GetScalarTable(density_element, integration_rule);
                data.density_basis.SetSize(0, 0);
            } else {
                data.density_reference.reset();
                data.density_basis.SetSize(quadrature_point_count, density_dof_count);
            }
            if (potential_element.GetMapType() == mfem::FiniteElement::VALUE) {
                data.potential_reference =
                    m_fem.GetReferenceTables().GetScalarTable(potential_element, integration_rule);
                data.potential_basis.SetSize(0, 0);
            } else {
                data.potential_reference.reset();
                data.potential_basis.SetSize(quadrature_point_count, potential_dof_count);
            }

            const int dimension = m_fem.mesh->Dimension();
            if (mode == PreparationMode::linearization) {
                data.inverse_element_jacobians.SetSize(quadrature_point_count, dimension * dimension);
                data.displacement_reference = m_fem.GetReferenceTables().GetScalarTable(
                    *m_fem.displacementFes->GetFE(element_id), integration_rule
                );
            }

            data.quadrature_data.SetSize(quadrature_point_count);

            mfem::Vector density_shape;
            mfem::Vector potential_shape;
            if (!data.density_reference) {
                density_shape.SetSize(density_dof_count);
            }
            if (!data.potential_reference) {
                potential_shape.SetSize(potential_dof_count);
            }

            for (int quadrature_point = 0; quadrature_point < quadrature_point_count; ++quadrature_point) {
                const mfem::IntegrationPoint &integration_point = integration_rule.IntPoint(quadrature_point);

                transformation.SetIntPoint(&integration_point);

                // VALUE maps use the shared reference basis. Preserve the
                // physical-shape evaluation for every other scalar map type.
                if (!data.density_reference) {
                    density_element.CalcPhysShape(transformation, density_shape);
                    for (int i = 0; i < density_dof_count; ++i) {
                        data.density_basis(quadrature_point, i) = density_shape(i);
                    }
                }
                if (!data.potential_reference) {
                    potential_element.CalcPhysShape(transformation, potential_shape);
                    for (int i = 0; i < potential_dof_count; ++i) {
                        data.potential_basis(quadrature_point, i) = potential_shape(i);
                    }
                }

                const double coefficient_value = source_coefficient.Eval(transformation, integration_point);

                if (source_coefficient.GetMappingFailure() != mapping::MappingStatus::valid ||
                    source_coefficient.HasNonFiniteArithmetic()) {
                    break;
                }

                if (mode == PreparationMode::linearization) {
                    const mfem::DenseMatrix &inverse_element_jacobian = source_coefficient.GetInverseElementJacobian();
                    for (int row = 0; row < dimension; ++row) {
                        for (int column = 0; column < dimension; ++column) {
                            data.inverse_element_jacobians(quadrature_point, row * dimension + column) =
                                inverse_element_jacobian(row, column);
                        }
                    }
                }

                transformation.SetIntPoint(&integration_point);

                const double quadrature_value = integration_point.weight * transformation.Weight() * coefficient_value;

                if (!std::isfinite(quadrature_value) || quadrature_value <= 0.0) {
                    localNonFiniteQuadrature = true;
                    break;
                }

                data.quadrature_data(quadrature_point) = quadrature_value;
            }

            if (source_coefficient.GetMappingFailure() != mapping::MappingStatus::valid ||
                source_coefficient.HasNonFiniteArithmetic() || localNonFiniteQuadrature) {
                break;
            }
        }
        m_elements.resize(prepared_element_count);

        const bool localNonFiniteArithmetic = source_coefficient.HasNonFiniteArithmetic() || localNonFiniteQuadrature;
        auto preparationResult              = synchronize_preparation_failure(
            source_coefficient.GetMappingFailure(), localNonFiniteArithmetic, m_fem.mesh->GetComm()
        );
        if (!preparationResult.has_value()) {
            return preparationResult;
        }

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

        m_is_prepared        = true;
        m_has_variation_data = mode == PreparationMode::linearization;
        ++m_preparation_count;
        return {};
    }
    void PreparedMappedGravitySourceOperator::Mult(
        const mfem::Vector &density,
        mfem::Vector &action
    ) const {
        MEAN_FIELD_PROFILE_SCOPE("PreparedMappedGravitySourceOperator::Mult");

        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedGravitySourceOperator must be prepared before "
                           "Mult is called."
        );

        MFEM_VERIFY(
            density.Size() == Width(), "PreparedMappedGravitySourceOperator received a density vector "
                                       "with the wrong size."
        );

        m_density_true.SetSize(m_density_map.full_size());
        m_density_map.scatter(density, m_density_true);

        true_to_local(*m_fem.densityFes, m_density_true, m_density_local);

        m_local_action.SetSize(m_fem.gravityPotentialFes->GetVSize());
        m_local_action = 0.0;

        for (const ElementPAData &data : m_elements) {
            m_density_local.GetSubVector(data.density_dofs, m_element_input);

            if (data.density_dof_transformation != nullptr) {
                data.density_dof_transformation->InvTransformPrimal(m_element_input);
            }

            m_quadrature_action.SetSize(data.quadrature_data.Size());

            // B_density * x_e
            data.GetDensityBasis().Mult(m_element_input, m_quadrature_action);

            // D * B_density * x_e
            for (int q = 0; q < m_quadrature_action.Size(); ++q) {
                m_quadrature_action(q) *= data.quadrature_data(q);
            }

            m_element_action.SetSize(data.potential_dofs.Size());

            // B_potential^T * D * B_density * x_e
            data.GetPotentialBasis().MultTranspose(m_quadrature_action, m_element_action);

            if (data.potential_dof_transformation != nullptr) {
                data.potential_dof_transformation->TransformDual(m_element_action);
            }

            m_local_action.AddElementVector(data.potential_dofs, m_element_action);
        }

        if (m_potential_map.is_identity()) {
            local_to_true(*m_fem.gravityPotentialFes, m_local_action, action);
        } else {
            local_to_true(*m_fem.gravityPotentialFes, m_local_action, m_action_true);
            action.SetSize(Height());
            m_potential_map.gather(m_action_true, action);
        }
    }

    void PreparedMappedGravitySourceOperator::MultDisplacementVariationTrue(
        const mfem::Vector &densityTrue,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &actionVariationTrue
    ) const {
        MFEM_VERIFY(
            m_is_prepared,
            "PreparedMappedGravitySourceOperator must be prepared before applying a displacement variation."
        );
        MFEM_VERIFY(
            m_has_variation_data,
            "PreparedMappedGravitySourceOperator requires linearization preparation before applying a displacement "
            "variation."
        );
        MFEM_VERIFY(
            densityTrue.Size() == m_fem.densityFes->GetTrueVSize(), "The full density 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.densityFes, densityTrue, m_density_local);
        true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacement_variation_local);

        m_local_variation_action.SetSize(m_fem.gravityPotentialFes->GetVSize());
        m_local_variation_action = 0.0;

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

        for (const ElementPAData &data : m_elements) {
            MFEM_VERIFY(
                data.integration_rule != nullptr,
                "Prepared gravity source displacement variation has no integration rule."
            );

            m_density_local.GetSubVector(data.density_dofs, m_element_density);
            m_displacement_variation_local.GetSubVector(data.displacement_dofs, m_element_displacement_variation);

            if (data.density_dof_transformation != nullptr) {
                data.density_dof_transformation->InvTransformPrimal(m_element_density);
            }
            if (data.displacement_dof_transformation != nullptr) {
                data.displacement_dof_transformation->InvTransformPrimal(m_element_displacement_variation);
            }

            const mfem::FiniteElement &displacement_element       = *m_fem.displacementFes->GetFE(data.element_id);
            const mapping::ElementDisplacementData direction_data = mapping::ElementDisplacementDataFromElementVDofs(
                displacement_element, m_element_displacement_variation
            );
            const mfem::DenseMatrix &direction_dofs = direction_data.GetDofMatrix();

            MFEM_VERIFY(
                data.inverse_element_jacobians.Height() == data.integration_rule->GetNPoints() &&
                    data.inverse_element_jacobians.Width() == dimension * dimension,
                "Prepared gravity source inverse-Jacobian data has an incompatible size."
            );

            m_reference_displacement_jacobian.SetSize(dimension, dimension);
            m_quadrature_variation_action.SetSize(data.integration_rule->GetNPoints());
            data.GetDensityBasis().Mult(m_element_density, m_quadrature_variation_action);

            for (int quadrature_point = 0; quadrature_point < data.integration_rule->GetNPoints(); ++quadrature_point) {
                mfem::MultAtB(
                    direction_dofs, data.displacement_reference->GetGradients(quadrature_point),
                    m_reference_displacement_jacobian
                );

                double logarithmic_jacobian_variation{0.0};
                for (int row = 0; row < dimension; ++row) {
                    for (int column = 0; column < dimension; ++column) {
                        logarithmic_jacobian_variation +=
                            data.inverse_element_jacobians(quadrature_point, row * dimension + column) *
                            m_reference_displacement_jacobian(column, row);
                    }
                }

                m_quadrature_variation_action(quadrature_point) *=
                    data.quadrature_data(quadrature_point) * logarithmic_jacobian_variation;
                MFEM_VERIFY(
                    std::isfinite(m_quadrature_variation_action(quadrature_point)),
                    "Prepared gravity source displacement variation encountered a non-finite quadrature value."
                );
            }

            m_element_variation_action.SetSize(data.potential_dofs.Size());
            data.GetPotentialBasis().MultTranspose(m_quadrature_variation_action, m_element_variation_action);

            if (data.potential_dof_transformation != nullptr) {
                data.potential_dof_transformation->TransformDual(m_element_variation_action);
            }
            m_local_variation_action.AddElementVector(data.potential_dofs, m_element_variation_action);
        }

        local_to_true(*m_fem.gravityPotentialFes, m_local_variation_action, actionVariationTrue);
    }

    void PreparedMappedGravitySourceOperator::MultTranspose(
        const mfem::Vector &potential,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedGravitySourceOperator must be prepared before "
                           "MultTranspose is called."
        );

        MFEM_VERIFY(
            potential.Size() == Height(), "PreparedMappedGravitySourceOperator received a potential vector "
                                          "with the wrong size."
        );

        if (m_potential_map.is_identity()) {
            true_to_local(*m_fem.gravityPotentialFes, potential, m_potential_local);
        } else {
            m_potential_true.SetSize(m_potential_map.full_size());
            m_potential_map.scatter(potential, m_potential_true);
            true_to_local(*m_fem.gravityPotentialFes, m_potential_true, m_potential_local);
        }

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

        for (const ElementPAData &data : m_elements) {
            m_potential_local.GetSubVector(data.potential_dofs, m_element_input);

            if (data.potential_dof_transformation != nullptr) {
                data.potential_dof_transformation->InvTransformPrimal(m_element_input);
            }

            m_quadrature_action.SetSize(data.quadrature_data.Size());

            data.GetPotentialBasis().Mult(m_element_input, m_quadrature_action);

            for (int q = 0; q < m_quadrature_action.Size(); ++q) {
                m_quadrature_action(q) *= data.quadrature_data(q);
            }

            m_element_action.SetSize(data.density_dofs.Size());

            data.GetDensityBasis().MultTranspose(m_quadrature_action, m_element_action);

            if (data.density_dof_transformation != nullptr) {
                data.density_dof_transformation->TransformDual(m_element_action);
            }

            m_local_action.AddElementVector(data.density_dofs, m_element_action);
        }

        local_to_true(*m_fem.densityFes, m_local_action, m_action_true);
        action.SetSize(Width());
        m_density_map.gather(m_action_true, action);
    }
    bool PreparedMappedGravitySourceOperator::IsPrepared() const noexcept {
        return m_is_prepared;
    }

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

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

    const field::FieldDofMap &PreparedMappedGravitySourceOperator::GetDensityMap() const noexcept {
        return m_density_map;
    }

    const field::FieldDofMap &PreparedMappedGravitySourceOperator::GetPotentialMap() const noexcept {
        return m_potential_map;
    }

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