module;
#include <cmath>
#include <mfem.hpp>

module mean_field;
import :operators.kernels.gravity_field;

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

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

    void true_to_local(
        const mfem::ParFiniteElementSpace &finite_element_space,
        const mfem::Vector &true_vector,
        mfem::Vector &local_vector
    ) {
        MFEM_VERIFY(true_vector.Size() == finite_element_space.GetTrueVSize(), "True vector has the wrong size.");

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

    void add_local_to_true(
        const mfem::ParFiniteElementSpace &fes,
        const mfem::Vector &local_vector,
        mfem::Vector &true_vector
    ) {
        const mfem::Operator *prolongation = fes.GetProlongationMatrix();
        if (prolongation != nullptr) {
            prolongation->AddMultTranspose(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>;
        MFEM_VERIFY(
            element.GetOrder() == mean_field::field::Gravity::Flux::familyOrder + 1,
            "The H(div) kernel element does not match the registered gravity "
            "flux."
        );
        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;
    }

    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 source-kernel trial element does not match the registered "
            "density "
            "field."
        );
        MFEM_VERIFY(
            potential_element.GetOrder() == mean_field::field::Gravity::Potential::familyOrder,
            "The source-kernel 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
            );

        const auto resolution = f.quadratureFactory->get(query, transformation.GetGeometryType());
        MFEM_VERIFY(
            resolution.integration_rule != nullptr, "The quadrature policy did not return a gravity-source integration "
                                                    "rule."
        );
        return *resolution.integration_rule;
    }
} // namespace

namespace mean_field::operators::kernels {
    void apply_mapped_hdiv_mass(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper,
        const mfem::Vector &gravity_gradient_true,
        const mfem::Vector &displacement_true,
        mfem::Vector &action
    ) {
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "The H(div) mass kernel requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "The H(div) mass kernel requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "The H(div) mass kernel requires the compactification "
                                              "finite-element "
                                              "space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "The H(div) mass kernel requires the compactification field."
        );
        MFEM_VERIFY(f.quadratureFactory != nullptr, "The H(div) mass kernel requires the quadrature rule factory.");
        MFEM_VERIFY(
            gravity_gradient_true.Size() == f.gravityFluxFes->GetTrueVSize(),
            "The gravity-gradient vector has the wrong size."
        );
        MFEM_VERIFY(
            displacement_true.Size() == f.displacementFes->GetTrueVSize(), "The displacement vector has the wrong size."
        );

        mfem::Vector gravity_gradient_local;
        mfem::Vector displacement_local;
        true_to_local(*f.gravityFluxFes, gravity_gradient_true, gravity_gradient_local);
        true_to_local(*f.displacementFes, displacement_true, displacement_local);

        mfem::Vector local_action(f.gravityFluxFes->GetVSize());
        local_action = 0.0;

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

        mfem::Array<int> gravity_dofs;
        mfem::Array<int> displacement_dofs;
        mfem::Array<int> compactification_dofs;

        mfem::Vector element_gravity_gradient;
        mfem::Vector element_displacement;
        mfem::Vector element_compactification;
        mfem::Vector element_action;
        mfem::Vector gravity_gradient_value;
        mfem::Vector mapped_gravity_gradient_value;

        mfem::DenseMatrix vector_shape;
        mfem::DenseMatrix mapped_mass_tensor;

        for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
            const mfem::FiniteElement &gravity_element          = *f.gravityFluxFes->GetFE(element_id);
            const mfem::FiniteElement &displacement_element     = *f.displacementFes->GetFE(element_id);
            const mfem::FiniteElement &compactification_element = *f.compactificationFes->GetFE(element_id);
            mfem::ElementTransformation *transformation         = f.mesh->GetElementTransformation(element_id);

            mfem::DofTransformation *gravity_dof_transformation =
                f.gravityFluxFes->GetElementVDofs(element_id, gravity_dofs);
            mfem::DofTransformation *displacement_dof_transformation =
                f.displacementFes->GetElementVDofs(element_id, displacement_dofs);
            mfem::DofTransformation *compactification_dof_transformation =
                f.compactificationFes->GetElementDofs(element_id, compactification_dofs);

            gravity_gradient_local.GetSubVector(gravity_dofs, element_gravity_gradient);
            displacement_local.GetSubVector(displacement_dofs, element_displacement);
            f.compactificationCoordinate->GetSubVector(compactification_dofs, element_compactification);

            if (gravity_dof_transformation != nullptr)
                gravity_dof_transformation->InvTransformPrimal(element_gravity_gradient);
            if (displacement_dof_transformation != nullptr)
                displacement_dof_transformation->InvTransformPrimal(element_displacement);
            if (compactification_dof_transformation != nullptr)
                compactification_dof_transformation->InvTransformPrimal(element_compactification);

            // const mapping::ElementDisplacementData
            // displacement_data(displacement_element, element_displacement,
            // mfem::Ordering::byVDIM);

            const mapping::ElementDisplacementData displacement_data =
                mapping::ElementDisplacementDataFromElementVDofs(displacement_element, element_displacement);

            const mapping::ElementCompactificationData compactification_data(
                compactification_element, element_compactification
            );
            const mapping::ElementMappingData mapping_data{
                .displacement = displacement_data, .compactification = compactification_data
            };

            const int gravity_dof_count = gravity_element.GetDof();
            const int dimension         = transformation->GetSpaceDim();

            element_action.SetSize(gravity_dof_count);
            gravity_gradient_value.SetSize(dimension);
            mapped_gravity_gradient_value.SetSize(dimension);
            vector_shape.SetSize(gravity_dof_count, dimension);
            mapped_mass_tensor.SetSize(dimension);
            element_action = 0.0;

            const mfem::IntegrationRule &integration_rule =
                get_hdiv_mass_rule(f, domain_mapper, gravity_element, *transformation);

            for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
                const mfem::IntegrationPoint &integration_point = integration_rule.IntPoint(q);
                transformation->SetIntPoint(&integration_point);

                mapping::VolumeMappingContext mapping_context;
                const mapping::MappingStatus status = domain_mapper.EvaluateVolume(
                    mapping_data, *transformation, integration_point, workspace, mapping_context
                );
                MFEM_VERIFY(
                    status == mapping::MappingStatus::valid,
                    "Stateless mapping failed in the matrix-free H(div) mass "
                    "kernel. "
                    "Element: "
                        << element_id << ", attribute: " << transformation->Attribute << ", quadrature point: " << q
                        << ", status: " << static_cast<int>(status)
                );

                gravity_element.CalcVShape(*transformation, vector_shape);
                mapping::ComputeHDivMassTensor(mapping_context.mapping, mapped_mass_tensor);

                vector_shape.MultTranspose(element_gravity_gradient, gravity_gradient_value);
                mapped_mass_tensor.Mult(gravity_gradient_value, mapped_gravity_gradient_value);

                const double weight = integration_point.weight * transformation->Weight();

                for (int i = 0; i < gravity_dof_count; ++i) {
                    double value = 0.0;
                    for (int component = 0; component < dimension; ++component)
                        value += vector_shape(i, component) * mapped_gravity_gradient_value(component);
                    element_action(i) += weight * value;
                }
            }

            if (gravity_dof_transformation != nullptr)
                gravity_dof_transformation->TransformDual(element_action);
            local_action.AddElementVector(gravity_dofs, element_action);
        }

        local_to_true(*f.gravityFluxFes, local_action, action);
    }

    void apply_mapped_source(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper,
        const mfem::Vector &density_true,
        const mfem::Vector &displacement_true,
        mfem::Vector &action
    ) {
        MFEM_VERIFY(
            f.densityFes != nullptr, "The gravity-source kernel requires the "
                                     "density finite-element space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "The gravity-source kernel requires the gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "The gravity-source kernel requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "The gravity-source kernel requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "The gravity-source kernel requires the compactification field."
        );
        MFEM_VERIFY(f.quadratureFactory != nullptr, "The gravity-source kernel requires the quadrature rule factory.");
        MFEM_VERIFY(density_true.Size() == f.densityFes->GetTrueVSize(), "The density vector has the wrong size.");
        MFEM_VERIFY(
            displacement_true.Size() == f.displacementFes->GetTrueVSize(), "The displacement vector has the wrong size."
        );

        mfem::Vector density_local;
        mfem::Vector displacement_local;
        true_to_local(*f.densityFes, density_true, density_local);
        true_to_local(*f.displacementFes, displacement_true, displacement_local);

        mfem::Vector local_action(f.gravityPotentialFes->GetVSize());
        local_action = 0.0;

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

        mfem::Array<int> density_dofs;
        mfem::Array<int> potential_dofs;
        mfem::Array<int> displacement_dofs;
        mfem::Array<int> compactification_dofs;

        mfem::Vector element_density;
        mfem::Vector element_displacement;
        mfem::Vector element_compactification;
        mfem::Vector element_action;
        mfem::Vector density_shape;
        mfem::Vector potential_shape;

        constexpr double source_scale = 4.0 * M_PI * utils::G;

        for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
            mfem::ElementTransformation *transformation = f.mesh->GetElementTransformation(element_id);
            if (is_vacuum_attribute(transformation->Attribute))
                continue;

            const mfem::FiniteElement &density_element          = *f.densityFes->GetFE(element_id);
            const mfem::FiniteElement &potential_element        = *f.gravityPotentialFes->GetFE(element_id);
            const mfem::FiniteElement &displacement_element     = *f.displacementFes->GetFE(element_id);
            const mfem::FiniteElement &compactification_element = *f.compactificationFes->GetFE(element_id);

            mfem::DofTransformation *density_dof_transformation =
                f.densityFes->GetElementDofs(element_id, density_dofs);
            mfem::DofTransformation *potential_dof_transformation =
                f.gravityPotentialFes->GetElementDofs(element_id, potential_dofs);
            mfem::DofTransformation *displacement_dof_transformation =
                f.displacementFes->GetElementVDofs(element_id, displacement_dofs);
            mfem::DofTransformation *compactification_dof_transformation =
                f.compactificationFes->GetElementDofs(element_id, compactification_dofs);

            density_local.GetSubVector(density_dofs, element_density);
            displacement_local.GetSubVector(displacement_dofs, element_displacement);
            f.compactificationCoordinate->GetSubVector(compactification_dofs, element_compactification);

            if (density_dof_transformation != nullptr)
                density_dof_transformation->InvTransformPrimal(element_density);
            if (displacement_dof_transformation != nullptr)
                displacement_dof_transformation->InvTransformPrimal(element_displacement);
            if (compactification_dof_transformation != nullptr)
                compactification_dof_transformation->InvTransformPrimal(element_compactification);

            const mapping::ElementDisplacementData displacement_data =
                mapping::ElementDisplacementDataFromElementVDofs(displacement_element, element_displacement);
            const mapping::ElementCompactificationData compactification_data(
                compactification_element, element_compactification
            );
            const mapping::ElementMappingData mapping_data{
                .displacement = displacement_data, .compactification = compactification_data
            };

            const int density_dof_count   = density_element.GetDof();
            const int potential_dof_count = potential_element.GetDof();

            density_shape.SetSize(density_dof_count);
            potential_shape.SetSize(potential_dof_count);
            element_action.SetSize(potential_dof_count);
            element_action = 0.0;

            const mfem::IntegrationRule &integration_rule =
                get_source_rule(f, density_element, potential_element, *transformation);

            for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
                const mfem::IntegrationPoint &integration_point = integration_rule.IntPoint(q);
                transformation->SetIntPoint(&integration_point);

                mapping::VolumeMappingContext mapping_context;
                const mapping::MappingStatus status = domain_mapper.EvaluateVolume(
                    mapping_data, *transformation, integration_point, workspace, mapping_context
                );
                MFEM_VERIFY(
                    status == mapping::MappingStatus::valid,
                    "Stateless mapping failed in the matrix-free "
                    "gravity-source "
                    "kernel. Element: "
                        << element_id << ", attribute: " << transformation->Attribute << ", quadrature point: " << q
                        << ", status: " << static_cast<int>(status)
                );

                density_element.CalcShape(integration_point, density_shape);
                potential_element.CalcShape(integration_point, potential_shape);

                const double density_value = element_density * density_shape;
                const double weight        = source_scale * density_value * mapping_context.quadrature.weight;

                for (int i = 0; i < potential_dof_count; ++i)
                    element_action(i) += weight * potential_shape(i);
            }

            if (potential_dof_transformation != nullptr)
                potential_dof_transformation->TransformDual(element_action);
            local_action.AddElementVector(potential_dofs, element_action);
        }

        local_to_true(*f.gravityPotentialFes, local_action, action);
    }

    void apply_mapped_hdiv_mass_variation(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper,
        const mfem::Vector &gravity_gradient_true,
        const mfem::Vector &displacement_true,
        const mfem::Vector &displacement_variation_true,
        mfem::Vector &action_variation
    ) {
        MFEM_VERIFY(f.mesh != nullptr, "The H(div) mass-variation kernel requires a mesh.");
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "The H(div) mass-variation kernel requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "The H(div) mass-variation kernel requires "
                                          "the displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "The H(div) mass-variation kernel requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "The H(div) mass-variation kernel requires the compactification "
                                                     "field."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "The H(div) mass-variation kernel requires the quadrature rule "
                                            "factory."
        );
        MFEM_VERIFY(
            gravity_gradient_true.Size() == f.gravityFluxFes->GetTrueVSize(),
            "The gravity-gradient vector has the wrong size."
        );
        MFEM_VERIFY(
            displacement_true.Size() == f.displacementFes->GetTrueVSize(), "The displacement vector has the wrong size."
        );
        MFEM_VERIFY(
            displacement_variation_true.Size() == f.displacementFes->GetTrueVSize(),
            "The displacement-variation vector has the wrong size."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "The domain-mapper dimension does not match the mesh dimension."
        );

        mfem::Vector gravity_gradient_local;
        mfem::Vector displacement_local;
        mfem::Vector displacement_variation_local;

        true_to_local(*f.gravityFluxFes, gravity_gradient_true, gravity_gradient_local);
        true_to_local(*f.displacementFes, displacement_true, displacement_local);
        true_to_local(*f.displacementFes, displacement_variation_true, displacement_variation_local);

        mfem::Vector local_action(f.gravityFluxFes->GetVSize());
        local_action = 0.0;

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

        mfem::Array<int> gravity_gradient_dofs;
        mfem::Array<int> displacement_dofs;
        mfem::Array<int> compactification_dofs;

        mfem::Vector element_gravity_gradient;
        mfem::Vector element_displacement;
        mfem::Vector element_displacement_variation;
        mfem::Vector element_compactification;
        mfem::Vector element_action;
        mfem::Vector gravity_gradient_value;
        mfem::Vector mass_tensor_variation_action;

        mfem::DenseMatrix gravity_gradient_shape;
        mfem::DenseMatrix mass_tensor_variation;
        mapping::VolumeMappingContext mapping_context;
        mapping::VolumeMappingVariation mapping_variation;

        for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
            const mfem::FiniteElement &gravity_gradient_element = *f.gravityFluxFes->GetFE(element_id);
            const mfem::FiniteElement &displacement_element     = *f.displacementFes->GetFE(element_id);
            const mfem::FiniteElement &compactification_element = *f.compactificationFes->GetFE(element_id);
            mfem::ElementTransformation *transformation         = f.mesh->GetElementTransformation(element_id);

            MFEM_VERIFY(
                transformation != nullptr, "The H(div) mass-variation kernel "
                                           "received a null element transformation."
            );

            mfem::DofTransformation *gravity_dof_transformation =
                f.gravityFluxFes->GetElementVDofs(element_id, gravity_gradient_dofs);
            mfem::DofTransformation *displacement_dof_transformation =
                f.displacementFes->GetElementVDofs(element_id, displacement_dofs);
            mfem::DofTransformation *compactification_dof_transformation =
                f.compactificationFes->GetElementDofs(element_id, compactification_dofs);

            gravity_gradient_local.GetSubVector(gravity_gradient_dofs, element_gravity_gradient);
            displacement_local.GetSubVector(displacement_dofs, element_displacement);
            displacement_variation_local.GetSubVector(displacement_dofs, element_displacement_variation);
            f.compactificationCoordinate->GetSubVector(compactification_dofs, element_compactification);

            if (gravity_dof_transformation != nullptr)
                gravity_dof_transformation->InvTransformPrimal(element_gravity_gradient);

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

            if (compactification_dof_transformation != nullptr)
                compactification_dof_transformation->InvTransformPrimal(element_compactification);

            const mapping::ElementDisplacementData displacement_data =
                mapping::ElementDisplacementDataFromElementVDofs(displacement_element, element_displacement);
            const mapping::ElementDisplacementData displacement_variation_data =
                mapping::ElementDisplacementDataFromElementVDofs(displacement_element, element_displacement_variation);
            const mapping::ElementCompactificationData compactification_data(
                compactification_element, element_compactification
            );
            const mapping::ElementMappingData mapping_data{
                .displacement = displacement_data, .compactification = compactification_data
            };

            const int gravity_gradient_dof_count = gravity_gradient_element.GetDof();
            const int dimension                  = transformation->GetSpaceDim();

            element_action.SetSize(gravity_gradient_dof_count);
            element_action = 0.0;

            gravity_gradient_value.SetSize(dimension);
            mass_tensor_variation_action.SetSize(dimension);

            gravity_gradient_shape.SetSize(gravity_gradient_dof_count, dimension);
            mass_tensor_variation.SetSize(dimension, dimension);

            const mfem::IntegrationRule &integration_rule =
                get_hdiv_mass_rule(f, domain_mapper, gravity_gradient_element, *transformation);

            for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
                const mfem::IntegrationPoint &integration_point = integration_rule.IntPoint(q);
                transformation->SetIntPoint(&integration_point);

                const mapping::MappingStatus status = domain_mapper.EvaluateVolume(
                    mapping_data, *transformation, integration_point, workspace, mapping_context
                );
                MFEM_VERIFY(
                    status == mapping::MappingStatus::valid,
                    "Stateless mapping failed in the matrix-free H(div) mass "
                    "kernel. "
                    "Element: "
                        << element_id << ", attribute: " << transformation->Attribute << ", quadrature point: " << q
                        << ", status: " << static_cast<int>(status)
                );

                const mapping::MappingStatus variation_status = domain_mapper.EvaluateVolumeVariation(
                    mapping_data, displacement_variation_data, *transformation, integration_point, mapping_context,
                    workspace, mapping_variation
                );
                MFEM_VERIFY(
                    variation_status == mapping::MappingStatus::valid,
                    "The mapping variation is invalid while applying the "
                    "H(div) mass "
                    "variation."
                );

                mapping::ComputeHDivMassTensorVariation(
                    mapping_context.mapping, mapping_variation.mapping, mass_tensor_variation
                );

                gravity_gradient_element.CalcVShape(*transformation, gravity_gradient_shape);
                gravity_gradient_shape.MultTranspose(element_gravity_gradient, gravity_gradient_value);
                mass_tensor_variation.Mult(gravity_gradient_value, mass_tensor_variation_action);
                const double reference_weight = integration_point.weight * transformation->Weight();

                gravity_gradient_shape.AddMult(mass_tensor_variation_action, element_action, reference_weight);
            }

            if (gravity_dof_transformation != nullptr)
                gravity_dof_transformation->TransformDual(element_action);
            local_action.AddElementVector(gravity_gradient_dofs, element_action);
        }

        action_variation.SetSize(f.gravityFluxFes->GetTrueVSize());
        action_variation = 0.0;
        add_local_to_true(*f.gravityFluxFes, local_action, action_variation);
    }

    void apply_mapped_source_variation(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper,
        const mfem::Vector &density_true,
        const mfem::Vector &displacement_true,
        const mfem::Vector &displacement_variation_true,
        mfem::Vector &action_variation
    ) {
        MFEM_VERIFY(f.mesh != nullptr, "The source-variation kernel requires a mesh.");
        MFEM_VERIFY(
            f.densityFes != nullptr, "The source-variation kernel requires the density finite-element "
                                     "space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "The source-variation kernel requires the gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "The source-variation kernel requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "The source-variation kernel requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "The source-variation kernel requires the compactification field."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "The source-variation kernel requires the quadrature rule factory."
        );
        MFEM_VERIFY(density_true.Size() == f.densityFes->GetTrueVSize(), "The density vector has the wrong size.");
        MFEM_VERIFY(
            displacement_true.Size() == f.displacementFes->GetTrueVSize(), "The displacement vector has the wrong size."
        );
        MFEM_VERIFY(
            displacement_variation_true.Size() == f.displacementFes->GetTrueVSize(),
            "The displacement-variation vector has the wrong size."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "The domain-mapper dimension does not match the mesh dimension."
        );

        mfem::Vector density_local;
        mfem::Vector displacement_local;
        mfem::Vector displacement_variation_local;

        true_to_local(*f.densityFes, density_true, density_local);
        true_to_local(*f.displacementFes, displacement_true, displacement_local);
        true_to_local(*f.displacementFes, displacement_variation_true, displacement_variation_local);

        mfem::Vector local_action(f.gravityPotentialFes->GetVSize());
        local_action = 0.0;

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

        mfem::Array<int> density_dofs;
        mfem::Array<int> potential_dofs;
        mfem::Array<int> displacement_dofs;
        mfem::Array<int> compactification_dofs;

        mfem::Vector element_density;
        mfem::Vector element_displacement;
        mfem::Vector element_displacement_variation;
        mfem::Vector element_compactification;
        mfem::Vector element_action;
        mfem::Vector density_shape;
        mfem::Vector potential_shape;

        mapping::VolumeMappingContext mapping_context;
        mapping::VolumeMappingVariation mapping_variation;

        constexpr double gravitational_source_scale = 4.0 * M_PI * utils::G;

        for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
            mfem::ElementTransformation *transformation = f.mesh->GetElementTransformation(element_id);
            MFEM_VERIFY(
                transformation != nullptr, "The source-variation kernel received a null element "
                                           "transformation."
            );

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

            const mfem::FiniteElement &density_element          = *f.densityFes->GetFE(element_id);
            const mfem::FiniteElement &potential_element        = *f.gravityPotentialFes->GetFE(element_id);
            const mfem::FiniteElement &displacement_element     = *f.displacementFes->GetFE(element_id);
            const mfem::FiniteElement &compactification_element = *f.compactificationFes->GetFE(element_id);

            mfem::DofTransformation *density_dof_transformation =
                f.densityFes->GetElementDofs(element_id, density_dofs);
            mfem::DofTransformation *potential_dof_transformation =
                f.gravityPotentialFes->GetElementDofs(element_id, potential_dofs);
            mfem::DofTransformation *displacement_dof_transformation =
                f.displacementFes->GetElementVDofs(element_id, displacement_dofs);
            mfem::DofTransformation *compactification_dof_transformation =
                f.compactificationFes->GetElementDofs(element_id, compactification_dofs);

            density_local.GetSubVector(density_dofs, element_density);
            displacement_local.GetSubVector(displacement_dofs, element_displacement);
            displacement_variation_local.GetSubVector(displacement_dofs, element_displacement_variation);
            f.compactificationCoordinate->GetSubVector(compactification_dofs, element_compactification);

            if (density_dof_transformation != nullptr)
                density_dof_transformation->InvTransformPrimal(element_density);

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

            if (compactification_dof_transformation != nullptr)
                compactification_dof_transformation->InvTransformPrimal(element_compactification);

            const mapping::ElementDisplacementData displacement_data =
                mapping::ElementDisplacementDataFromElementVDofs(displacement_element, element_displacement);
            const mapping::ElementDisplacementData displacement_variation_data =
                mapping::ElementDisplacementDataFromElementVDofs(displacement_element, element_displacement_variation);
            const mapping::ElementCompactificationData compactification_data(
                compactification_element, element_compactification
            );
            const mapping::ElementMappingData mapping_data{
                .displacement = displacement_data, .compactification = compactification_data
            };

            element_action.SetSize(potential_element.GetDof());
            element_action = 0.0;

            density_shape.SetSize(density_element.GetDof());
            potential_shape.SetSize(potential_element.GetDof());

            const mfem::IntegrationRule &integration_rule =
                get_source_rule(f, density_element, potential_element, *transformation);

            for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
                const mfem::IntegrationPoint &integration_point = integration_rule.IntPoint(q);
                transformation->SetIntPoint(&integration_point);

                const mapping::MappingStatus mapping_status = domain_mapper.EvaluateVolume(
                    mapping_data, *transformation, integration_point, workspace, mapping_context
                );
                MFEM_VERIFY(
                    mapping_status == mapping::MappingStatus::valid,
                    "The base mapping is invalid while applying the source "
                    "variation."
                );

                const mapping::MappingStatus variation_status = domain_mapper.EvaluateVolumeVariation(
                    mapping_data, displacement_variation_data, *transformation, integration_point, mapping_context,
                    workspace, mapping_variation
                );

                MFEM_VERIFY(
                    variation_status == mapping::MappingStatus::valid,
                    "The mapping variation is invalid while applying the "
                    "source "
                    "variation."
                );

                density_element.CalcShape(integration_point, density_shape);
                potential_element.CalcShape(integration_point, potential_shape);

                const double density_value = density_shape * element_density;
                const double source_variation_value =
                    gravitational_source_scale * density_value * mapping_variation.weight_variation;

                element_action.Add(source_variation_value, potential_shape);
            }

            if (potential_dof_transformation != nullptr)
                potential_dof_transformation->TransformDual(element_action);
            local_action.AddElementVector(potential_dofs, element_action);
        }

        action_variation.SetSize(f.gravityPotentialFes->GetTrueVSize());
        action_variation = 0.0;
        add_local_to_true(*f.gravityPotentialFes, local_action, action_variation);
    }
} // namespace mean_field::operators::kernels