module;

#include "interior_displacement_contractions.hpp"

#include <array>
#include <cmath>
#include <expected>
#include <mfem.hpp>
#include <optional>
#include <stdexcept>
#include <utility>

#include <mpi.h>

module mean_field;

import :operators.kernels.rotational_displacement_force;
import :operators.prepared_rotational_displacement_force;
import :fem.reference_tables;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using Rejection    = mean_field::operators::kernels::RotationalDisplacementForceRejection;
    using Reason       = mean_field::operators::kernels::RotationalDisplacementForceRejectionReason;

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

    [[nodiscard]] Rejection mapping_rejection(const mean_field::mapping::MappingStatus status) {
        MFEM_VERIFY(
            status != mean_field::mapping::MappingStatus::invalid_dimension,
            "Prepared rotational force mapping reported an invariant dimension mismatch."
        );
        return {.reason = Reason::invalid_mapping, .mappingStatus = status};
    }

    [[nodiscard]] Rejection non_finite_rejection() noexcept {
        return {.reason = Reason::non_finite_arithmetic};
    }

    [[nodiscard]] int encode_rejection(const std::optional<Rejection> &rejection) noexcept {
        if (!rejection.has_value()) {
            return 0;
        }
        if (rejection->reason == Reason::non_finite_arithmetic) {
            return 256;
        }
        return static_cast<int>(rejection->mappingStatus) + 1;
    }

    [[nodiscard]] Rejection decode_rejection(const int encoded) {
        if (encoded >= 256) {
            return non_finite_rejection();
        }
        return mapping_rejection(static_cast<mean_field::mapping::MappingStatus>(encoded - 1));
    }

    [[nodiscard]] std::expected<
        void,
        Rejection>
    synchronize_rejection(
        const std::optional<Rejection> &localRejection,
        const MPI_Comm communicator
    ) {
        const int localEncoded = encode_rejection(localRejection);
        int globalEncoded      = 0;
        if (MPI_Allreduce(&localEncoded, &globalEncoded, 1, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS) {
            throw std::runtime_error("Could not synchronize prepared rotational-force candidate validity.");
        }
        if (globalEncoded != 0) {
            return std::unexpected(decode_rejection(globalEncoded));
        }
        return {};
    }

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

    [[noreturn]] void throw_rejection(const Rejection &rejection) {
        if (rejection.reason == Reason::non_finite_arithmetic) {
            throw std::domain_error("Prepared rotational force produced non-finite arithmetic.");
        }
        throw std::domain_error("Prepared rotational force encountered an invalid mapped domain.");
    }

    void true_to_local(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &trueVector,
        mfem::Vector &localVector
    ) {
        localVector.SetSize(finiteElementSpace.GetVSize());
        const mfem::Operator *prolongation = finiteElementSpace.GetProlongationMatrix();
        if (prolongation != nullptr) {
            prolongation->Mult(trueVector, localVector);
        } else {
            localVector = trueVector;
        }
    }

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

    [[nodiscard]] int vector_dof_index(
        const mfem::Ordering::Type ordering,
        const int scalarDof,
        const int component,
        const int scalarDofCount,
        const int dimension
    ) {
        if (ordering == mfem::Ordering::byNODES) {
            return scalarDof + component * scalarDofCount;
        }
        MFEM_VERIFY(ordering == mfem::Ordering::byVDIM, "Unsupported displacement ordering.");
        return scalarDof * dimension + component;
    }

    [[nodiscard]] const mfem::IntegrationRule &get_rotation_force_rule(
        const mean_field::fem::FEM &f,
        const mfem::ElementTransformation &transformation
    ) {
        using DisplacementField = mean_field::field::Field<mean_field::field::Displacement>;
        const mean_field::quadrature::Query query =
            DisplacementField::make_query<mean_field::field::Displacement::Form::CentrifugalForce>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), std::array<int, 1>{1},
                mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
            );
        const mean_field::quadrature::MfemRule rule = f.quadratureFactory->get(query, transformation.GetGeometryType());
        MFEM_VERIFY(
            rule.integration_rule != nullptr,
            "The quadrature policy did not return a rotational-displacement-force integration rule."
        );
        return *rule.integration_rule;
    }
} // namespace

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

        MFEM_VERIFY(
            m_fem.mesh->Dimension() == 3, "PreparedRotationalDisplacementForceOperator requires a "
                                          "three-dimensional mesh."
        );

        MFEM_VERIFY(
            m_fem.densityFes != nullptr && m_fem.displacementFes != nullptr,
            "PreparedRotationalDisplacementForceOperator requires density "
            "and displacement finite-element spaces."
        );

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

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

        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_fem.mesh->Dimension(),
            "PreparedRotationalDisplacementForceOperator received a mapper "
            "with the wrong dimension."
        );
    }

    std::expected<
        void,
        kernels::RotationalDisplacementForceRejection>
    PreparedRotationalDisplacementForceOperator::TryPrepareElementData(const bool assembleResidual) {
        MFEM_VERIFY(m_rotation.has_value(), "Prepared rotational force has no frozen rotation state.");

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

        mfem::Vector baseDensityLocal;
        mfem::Vector baseDisplacementLocal;
        true_to_local(*m_fem.densityFes, m_context.GetBaseDensityTrue(), baseDensityLocal);
        true_to_local(*m_fem.displacementFes, m_context.GetDisplacementTrue(), baseDisplacementLocal);

        mapping::DomainMapper::Workspace workspace(m_domainMapper.GetDimension());
        mapping::VolumeMappingContext mappingContext;
        mfem::Array<int> compactificationDofs;
        mfem::Vector elementBaseDensity;
        mfem::Vector elementBaseDisplacement;
        mfem::Vector elementCompactification;
        mfem::Vector densityShape;
        mfem::Vector potentialGradient;

        if (assembleResidual) {
            m_localAction.SetSize(m_fem.displacementFes->GetVSize());
            m_localAction = 0.0;
        }
        mfem::Vector residualComponent;
        const int dimension = m_domainMapper.GetDimension();
        // Rigid rotation has a position-independent acceleration derivative.
        m_centrifugalJacobian.SetSize(dimension, dimension);
        mfem::Vector unit(dimension), derivative(dimension);
        for (int c = 0; c < dimension; ++c) {
            unit = 0.0; unit(c) = 1.0;
            m_rotation->potential_gradient_directional_derivative(unit, derivative);
            for (int r = 0; r < dimension; ++r) m_centrifugalJacobian(r, c) = -derivative(r);
        }

        for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
            mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(elementId);
            MFEM_VERIFY(transformation != nullptr, "Prepared rotational force received a null transformation.");
            if (is_vacuum_attribute(transformation->Attribute)) {
                continue;
            }

            m_elements.emplace_back();
            ElementPAData &data           = m_elements.back();
            data.elementId                = elementId;
            data.densityDofTransformation = m_fem.densityFes->GetElementDofs(elementId, data.densityDofs);
            data.displacementDofTransformation =
                m_fem.displacementFes->GetElementVDofs(elementId, data.displacementDofs);
            mfem::DofTransformation *compactificationDofTransformation =
                m_fem.compactificationFes->GetElementDofs(elementId, compactificationDofs);

            baseDensityLocal.GetSubVector(data.densityDofs, elementBaseDensity);
            baseDisplacementLocal.GetSubVector(data.displacementDofs, elementBaseDisplacement);
            m_fem.compactificationCoordinate->GetSubVector(compactificationDofs, elementCompactification);
            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(elementBaseDensity);
            }
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(elementBaseDisplacement);
            }
            if (compactificationDofTransformation != nullptr) {
                compactificationDofTransformation->InvTransformPrimal(elementCompactification);
            }

            const mfem::FiniteElement &densityElement          = *m_fem.densityFes->GetFE(elementId);
            const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(elementId);
            const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(elementId);
            data.integrationRule                               = &get_rotation_force_rule(m_fem, *transformation);
            data.densityReferenceTable = m_fem.GetReferenceTables().GetScalarTable(densityElement, *data.integrationRule);
            data.displacementReferenceTable = m_fem.GetReferenceTables().GetScalarTable(displacementElement, *data.integrationRule);
            data.fastApplication = dimension == 3 && m_fem.displacementFes->GetOrdering() == mfem::Ordering::byNODES &&
                displacementElement.GetMapType() == mfem::FiniteElement::VALUE &&
                displacementElement.GetDerivType() == mfem::FiniteElement::GRAD;

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

            const int quadraturePointCount = data.integrationRule->GetNPoints();
            data.inverseElementJacobians.SetSize(quadraturePointCount, dimension * dimension);
            data.centrifugalAccelerations.SetSize(quadraturePointCount, dimension);
            data.baseDensityValues.SetSize(quadraturePointCount);
            data.quadratureWeights.SetSize(quadraturePointCount);
            densityShape.SetSize(densityElement.GetDof());
            potentialGradient.SetSize(dimension);
            if (assembleResidual) m_quadratureForce.SetSize(quadraturePointCount, dimension);

            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
                const mapping::MappingStatus status            = m_domainMapper.EvaluateVolume(
                    mappingData, *transformation, integrationPoint, workspace, mappingContext
                );
                if (status != mapping::MappingStatus::valid) {
                    return std::unexpected(mapping_rejection(status));
                }
                MFEM_VERIFY(
                    !mappingContext.mapping.compactified,
                    "Prepared rotational force encountered compactification on a stellar element."
                );

                densityElement.CalcShape(integrationPoint, densityShape);
                m_rotation->potential_gradient(mappingContext.mapping.physical_position, potentialGradient);
                data.baseDensityValues(quadraturePoint) = elementBaseDensity * densityShape;
                data.quadratureWeights(quadraturePoint) = mappingContext.quadrature.weight;

                for (int row = 0; row < dimension; ++row) {
                    data.centrifugalAccelerations(quadraturePoint, row) = -potentialGradient(row);
                    for (int column = 0; column < dimension; ++column) {
                        data.inverseElementJacobians(quadraturePoint, row * dimension + column) =
                            mappingContext.quadrature.J_inv(row, column);
                    }
                }

                if (!std::isfinite(data.baseDensityValues(quadraturePoint)) ||
                    !std::isfinite(data.quadratureWeights(quadraturePoint)) || !vector_is_finite(potentialGradient)) {
                    return std::unexpected(non_finite_rejection());
                }
                for (int row = 0; row < dimension; ++row) {
                    for (int column = 0; column < dimension; ++column) {
                        if (!std::isfinite(data.inverseElementJacobians(quadraturePoint, row * dimension + column))) {
                            return std::unexpected(non_finite_rejection());
                        }
                    }
                }

                if (assembleResidual) {
                    const double weightedDensity =
                        data.quadratureWeights(quadraturePoint) * data.baseDensityValues(quadraturePoint);
                    for (int component = 0; component < dimension; ++component)
                        m_quadratureForce(quadraturePoint, component) = -weightedDensity * potentialGradient(component);
                }
            }

            if (assembleResidual) {
                const int ndof = displacementElement.GetDof();
                MFEM_VERIFY(data.displacementDofs.Size() == ndof * dimension,
                            "Prepared force residual has an invalid displacement element size.");
                m_elementAction.SetSize(data.displacementDofs.Size());
                residualComponent.SetSize(ndof);
                for (int component = 0; component < dimension; ++component) {
                    data.displacementReferenceTable->GetValues().MultTranspose(
                        m_quadratureForce.Data() + component * quadraturePointCount, residualComponent.GetData());
                    for (int dof = 0; dof < ndof; ++dof)
                        m_elementAction(vector_dof_index(m_fem.displacementFes->GetOrdering(),
                                                         dof, component, ndof, dimension)) = residualComponent(dof);
                }
                if (!vector_is_finite(m_elementAction)) return std::unexpected(non_finite_rejection());
                if (data.displacementDofTransformation != nullptr)
                    data.displacementDofTransformation->TransformDual(m_elementAction);
                m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
            }
        }
        if (assembleResidual && !vector_is_finite(m_localAction))
            return std::unexpected(non_finite_rejection());
        return {};
    }

    PreparedRotationalDisplacementForceReport PreparedRotationalDisplacementForceOperator::Prepare(
        const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
        const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        auto result = TryPrepare(state, dependencies, rotation);
        if (!result.has_value()) {
            throw_rejection(result.error());
        }
        return std::move(result).value();
    }

    std::expected<
        PreparedRotationalDisplacementForceReport,
        kernels::RotationalDisplacementForceRejection>
    PreparedRotationalDisplacementForceOperator::TryPrepare(
        const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
        const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        return TryPrepareImpl(state, dependencies, rotation, true);
    }

    PreparedRotationalDisplacementForceReport PreparedRotationalDisplacementForceOperator::PrepareReference(
        const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
        const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        auto result = TryPrepareImpl(state, dependencies, rotation, false);
        if (!result.has_value()) throw_rejection(result.error());
        return std::move(result).value();
    }

    std::expected<PreparedRotationalDisplacementForceReport, kernels::RotationalDisplacementForceRejection>
    PreparedRotationalDisplacementForceOperator::TryPrepareImpl(
        const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
        const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
        const physics::RigidRotation &rotation,
        const bool useFastPreparation
    ) {
        const bool wasPrepared = m_isPrepared;
        const bool implementationChanged = m_preparedWithFastResidual != useFastPreparation;
        const bool rotationChanged =
            !m_context.IsPrepared() || dependencies.rotation != m_context.GetDependencies().rotation;

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

        if (!report.contextReport.DidAnyWork() && wasPrepared && !implementationChanged) {
            return report;
        }

        m_isPrepared = false;

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

        MFEM_VERIFY(
            m_rotation.has_value(), "PreparedRotationalDisplacementForceOperator has no frozen "
                                    "rotation state."
        );

        if (report.contextReport.preparedBaseState || !wasPrepared || implementationChanged) {
            const bool inputsFinite = vector_is_finite(m_context.GetBaseDensityTrue()) &&
                                      vector_is_finite(m_context.GetDisplacementTrue());
            const auto inputStatus = synchronize_rejection(
                inputsFinite ? std::optional<Rejection>{} : std::optional<Rejection>{non_finite_rejection()},
                m_fem.mesh->GetComm());
            if (!inputStatus.has_value()) return std::unexpected(inputStatus.error());

            // Retain the original two-pass path for an independent residual
            // comparison. The production path evaluates the geometry once.
            if (!useFastPreparation) {
                auto residualResult = kernels::try_apply_rotational_displacement_force_residual(
                    m_fem, m_domainMapper, *m_rotation, m_context.GetBaseDensityTrue(), m_context.GetDisplacementTrue(),
                    m_actionTrue);
                if (!residualResult.has_value()) return std::unexpected(residualResult.error());
            }
            const auto elementResult = TryPrepareElementData(useFastPreparation);
            const std::optional<Rejection> localRejection = elementResult.has_value()
                ? std::optional<Rejection>{} : std::optional<Rejection>{elementResult.error()};
            const auto elementStatus = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
            if (!elementStatus.has_value()) return std::unexpected(elementStatus.error());

            if (useFastPreparation) local_to_true(*m_fem.displacementFes, m_localAction, m_actionTrue);
            m_cachedResidual.SetSize(m_context.GetDisplacementMap().reduced_size());
            m_context.GetDisplacementMap().gather(m_actionTrue, m_cachedResidual);
            const auto residualStatus = synchronize_rejection(
                vector_is_finite(m_cachedResidual) ? std::optional<Rejection>{}
                    : std::optional<Rejection>{non_finite_rejection()},
                m_fem.mesh->GetComm());
            if (!residualStatus.has_value()) return std::unexpected(residualStatus.error());

            m_preparedWithFastResidual = useFastPreparation;
            ++m_residualPreparationCount;
            report.preparedResidual = true;
        }

        MFEM_VERIFY(
            m_cachedResidual.Size() == m_context.GetDisplacementMap().reduced_size(),
            "The prepared rotational-displacement-force residual has the "
            "wrong size."
        );

        m_preparedDependencies = dependencies;
        m_isPrepared           = true;

        return report;
    }

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

    void PreparedRotationalDisplacementForceOperator::ApplyDensityJacobianAction(
        const mfem::Vector &densityVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_context.GetDensityMap().full_size());
        m_context.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);

        kernels::apply_rotational_displacement_force_density_action(
            m_fem, m_domainMapper, *m_rotation, m_densityVariationTrue, m_context.GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_densityJacobianStatistics.applications;
    }

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

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

        kernels::apply_rotational_displacement_force_displacement_action(
            m_fem, m_domainMapper, *m_rotation, m_context.GetBaseDensityTrue(), m_displacementVariationTrue,
            m_context.GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_displacementJacobianStatistics.applications;
    }

    void PreparedRotationalDisplacementForceOperator::ApplyPreparedCompleteJacobianActionTrue(
        const mfem::Vector &densityVariationTrue,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &actionTrue,
        const bool useFastPath, const bool useTensorPath
    ) const {
        true_to_local(*m_fem.densityFes, densityVariationTrue, m_densityVariationLocal);
        true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacementVariationLocal);
        m_localAction.SetSize(m_fem.displacementFes->GetVSize());
        m_localAction                       = 0.0;

        const int dimension                 = m_domainMapper.GetDimension();
        const mfem::Ordering::Type ordering = m_fem.displacementFes->GetOrdering();

        for (const ElementPAData &data : m_elements) {
            MFEM_VERIFY(data.integrationRule != nullptr, "Prepared rotational force has no integration rule.");

            m_densityVariationLocal.GetSubVector(data.densityDofs, m_elementDensityVariation);
            m_displacementVariationLocal.GetSubVector(data.displacementDofs, m_elementDisplacementVariation);
            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(m_elementDensityVariation);
            }
            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->InvTransformPrimal(m_elementDisplacementVariation);
            }

            const mfem::FiniteElement &densityElement      = *m_fem.densityFes->GetFE(data.elementId);
            const mfem::FiniteElement &displacementElement = *m_fem.displacementFes->GetFE(data.elementId);
            const mapping::ElementDisplacementData directionData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, m_elementDisplacementVariation);
            const mfem::DenseMatrix &directionDofs = directionData.GetDofMatrix();
            const int scalarDisplacementDofCount   = displacementElement.GetDof();

            if (useFastPath && data.fastApplication) {
                const int nq = data.integrationRule->GetNPoints();
                m_quadratureDensityVariation.SetSize(nq);
                if (useTensorPath && data.densityReferenceTable->HasTensorEvaluation3D())
                    data.densityReferenceTable->EvaluateValuesTensor3D(m_elementDensityVariation, m_quadratureDensityVariation, m_tensorWorkspace);
                else data.densityReferenceTable->GetValues().Mult(m_elementDensityVariation, m_quadratureDensityVariation);
                m_quadraturePositionVariation.SetSize(nq, 3);
                m_quadratureForce.SetSize(nq, 3);
                double *position = m_quadraturePositionVariation.Data();
                double *force = m_quadratureForce.Data();
                const double *u = directionDofs.Data();
                const auto &values = data.displacementReferenceTable->GetValues();
                const bool tensor = useTensorPath && data.displacementReferenceTable->HasTensorEvaluation3D();
                if (tensor) {
                    data.displacementReferenceTable->EvaluateTensor3D(directionDofs, m_tensorDisplacement, m_tensorWorkspace);
                    for (int c = 0; c < 3; ++c) for (int q = 0; q < nq; ++q)
                        position[q + c * nq] = m_tensorDisplacement(q, 4 * c);
                } else {
                    for (int c = 0; c < 3; ++c)
                        values.Mult(u + c * scalarDisplacementDofCount, position + c * nq);
                }
                for (int q = 0; q < nq; ++q) {
                    double trace = 0.0;
                    if (tensor) {
                        for (int row = 0; row < 3; ++row) for (int column = 0; column < 3; ++column)
                            trace += data.inverseElementJacobians(q, row * 3 + column) *
                                m_tensorDisplacement(q, 4 * column + row + 1);
                    } else trace = displacement_trace_3d(directionDofs,
                        data.displacementReferenceTable->GetGradients(q), data.inverseElementJacobians, q);
                    const double rho = data.baseDensityValues(q);
                    const double densityWeight = m_quadratureDensityVariation(q) + rho * trace;
                    for (int c = 0; c < 3; ++c) {
                        const double da = m_centrifugalJacobian(c, 0) * position[q] +
                            m_centrifugalJacobian(c, 1) * position[q + nq] +
                            m_centrifugalJacobian(c, 2) * position[q + 2 * nq];
                        force[q + c * nq] = data.quadratureWeights(q) *
                            (densityWeight * data.centrifugalAccelerations(q, c) + rho * da);
                    }
                }
                m_elementAction.SetSize(data.displacementDofs.Size());
                double *element = m_elementAction.HostWrite();
                if (tensor) {
                    data.displacementReferenceTable->ApplyTransposeTensor3D(m_quadratureForce, m_tensorElement, m_tensorWorkspace);
                    for (int c = 0; c < 3; ++c) for (int i = 0; i < scalarDisplacementDofCount; ++i)
                        element[i + c * scalarDisplacementDofCount] = m_tensorElement(i, c);
                } else {
                    for (int c = 0; c < 3; ++c)
                        values.MultTranspose(force + c * nq, element + c * scalarDisplacementDofCount);
                }
                if (data.displacementDofTransformation)
                    data.displacementDofTransformation->TransformDual(m_elementAction);
                m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
                continue;
            }

            m_densityShape.SetSize(densityElement.GetDof());
            m_displacementShape.SetSize(scalarDisplacementDofCount);
            m_referenceDisplacementDShape.SetSize(scalarDisplacementDofCount, dimension);
            m_referenceDisplacementJacobian.SetSize(dimension, dimension);
            m_physicalPositionVariation.SetSize(dimension);
            m_centrifugalAcceleration.SetSize(dimension);
            m_centrifugalAccelerationVariation.SetSize(dimension);
            m_weightedForce.SetSize(dimension);
            m_elementAction.SetSize(data.displacementDofs.Size());
            m_elementAction = 0.0;

            for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
                densityElement.CalcShape(integrationPoint, m_densityShape);
                displacementElement.CalcShape(integrationPoint, m_displacementShape);
                displacementElement.CalcDShape(integrationPoint, m_referenceDisplacementDShape);
                mfem::MultAtB(directionDofs, m_referenceDisplacementDShape, m_referenceDisplacementJacobian);
                directionDofs.MultTranspose(m_displacementShape, m_physicalPositionVariation);
                m_rotation->potential_gradient_directional_derivative(
                    m_physicalPositionVariation, m_centrifugalAccelerationVariation
                );
                m_centrifugalAccelerationVariation *= -1.0;

                double logarithmicJacobianVariation{0.0};
                for (int row = 0; row < dimension; ++row) {
                    m_centrifugalAcceleration(row) = data.centrifugalAccelerations(quadraturePoint, row);
                    for (int column = 0; column < dimension; ++column) {
                        logarithmicJacobianVariation +=
                            data.inverseElementJacobians(quadraturePoint, row * dimension + column) *
                            m_referenceDisplacementJacobian(column, row);
                    }
                }

                const double densityVariationValue = m_elementDensityVariation * m_densityShape;
                const double baseDensityValue      = data.baseDensityValues(quadraturePoint);
                m_weightedForce                    = 0.0;
                m_weightedForce.Add(densityVariationValue, m_centrifugalAcceleration);
                m_weightedForce.Add(baseDensityValue, m_centrifugalAccelerationVariation);
                m_weightedForce.Add(baseDensityValue * logarithmicJacobianVariation, m_centrifugalAcceleration);
                m_weightedForce *= data.quadratureWeights(quadraturePoint);

                for (int scalarDof = 0; scalarDof < scalarDisplacementDofCount; ++scalarDof) {
                    for (int component = 0; component < dimension; ++component) {
                        const int vectorDof =
                            vector_dof_index(ordering, scalarDof, component, scalarDisplacementDofCount, dimension);
                        m_elementAction(vectorDof) += m_displacementShape(scalarDof) * m_weightedForce(component);
                    }
                }
            }

            if (data.displacementDofTransformation != nullptr) {
                data.displacementDofTransformation->TransformDual(m_elementAction);
            }
            m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
        }

        local_to_true(*m_fem.displacementFes, m_localAction, actionTrue);
    }

    void PreparedRotationalDisplacementForceOperator::ApplyCompleteJacobianAction(
        const mfem::Vector &densityVariation,
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_context.GetDensityMap().full_size());
        m_displacementVariationTrue.SetSize(m_context.GetDisplacementMap().full_size());
        m_context.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_context.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        ApplyPreparedCompleteJacobianActionTrue(m_densityVariationTrue, m_displacementVariationTrue, m_actionTrue);
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_densityJacobianStatistics.applications;
        ++m_displacementJacobianStatistics.applications;
        ++m_completeJacobianStatistics.applications;
    }

    void PreparedRotationalDisplacementForceOperator::ApplyCompleteJacobianActionReference(
        const mfem::Vector &densityVariation,
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_context.GetDensityMap().full_size());
        m_displacementVariationTrue.SetSize(m_context.GetDisplacementMap().full_size());
        m_context.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_context.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        ApplyPreparedCompleteJacobianActionTrue(m_densityVariationTrue, m_displacementVariationTrue, m_actionTrue, false);
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

    }

    void PreparedRotationalDisplacementForceOperator::ApplyCompleteJacobianActionCachedReference(
        const mfem::Vector &densityVariation,
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_context.GetDensityMap().full_size());
        m_displacementVariationTrue.SetSize(m_context.GetDisplacementMap().full_size());
        m_context.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_context.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        ApplyPreparedCompleteJacobianActionTrue(m_densityVariationTrue, m_displacementVariationTrue, m_actionTrue, true, false);
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

    }

    int PreparedRotationalDisplacementForceOperator::GetTensorElementCount() const noexcept {
        int count = 0;
        for (const auto &data : m_elements)
            count += data.fastApplication && data.densityReferenceTable->HasTensorEvaluation3D() &&
                data.displacementReferenceTable->HasTensorEvaluation3D();
        return count;
    }

    bool PreparedRotationalDisplacementForceOperator::IsPrepared() const noexcept {
        return m_isPrepared && m_rotation.has_value() && m_context.MatchesDependencies(m_preparedDependencies);
    }

    const context::rotational_displacement_force::RotationalDisplacementForcePreparationStatistics &
    PreparedRotationalDisplacementForceOperator::GetContextPreparationStatistics() const noexcept {
        return m_context.GetPreparationStatistics();
    }

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

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

    const PreparedRotationalDisplacementForceColumnStatistics &
    PreparedRotationalDisplacementForceOperator::GetDensityJacobianStatistics() const noexcept {
        return m_densityJacobianStatistics;
    }

    const PreparedRotationalDisplacementForceColumnStatistics &
    PreparedRotationalDisplacementForceOperator::GetDisplacementJacobianStatistics() const noexcept {
        return m_displacementJacobianStatistics;
    }

    const PreparedRotationalDisplacementForceCompleteStatistics &
    PreparedRotationalDisplacementForceOperator::GetCompleteJacobianStatistics() const noexcept {
        return m_completeJacobianStatistics;
    }

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

    const context::rotational_displacement_force::RotationalDisplacementForceLinearizationContext &
    PreparedRotationalDisplacementForceOperator::GetContext() const noexcept {
        return m_context;
    }

    void PreparedRotationalDisplacementForceOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedRotationalDisplacementForceOperator must be prepared "
                          "for the current revisions before residual or Jacobian "
                          "application."
        );
    }

    PreparedRotationalDisplacementForceJacobianOperator::PreparedRotationalDisplacementForceJacobianOperator(
        const RotationalDisplacementForceLayout &layout,
        const PreparedRotationalDisplacementForceOperator &preparedOperator
    )
        : mfem::Operator(
              layout.residual_offsets().Last(),
              layout.value_offsets().Last()
          ),
          m_layout(layout),
          m_preparedOperator(preparedOperator) {
        using Form                  = utils::blocks::barotropic_equilibrium_form;

        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);

        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);

        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);

        MFEM_VERIFY(
            m_layout.size(densityValue) == m_preparedOperator.GetContext().GetDensityMap().reduced_size() &&
                m_layout.size(displacementValue) ==
                    m_preparedOperator.GetContext().GetDisplacementMap().reduced_size() &&
                m_layout.size(displacementResidual) ==
                    m_preparedOperator.GetContext().GetDisplacementMap().reduced_size(),
            "Prepared rotational-displacement-force MFEM adapter received "
            "incompatible coupled block sizes."
        );
    }

    void PreparedRotationalDisplacementForceJacobianOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_preparedOperator.IsPrepared(), "Prepared rotational-displacement-force MFEM adapter requires "
                                             "a prepared operator."
        );

        MFEM_VERIFY(
            direction.Size() == Width(), "Prepared rotational-displacement-force MFEM adapter received "
                                         "a direction with the wrong size."
        );

        using Form                  = utils::blocks::barotropic_equilibrium_form;

        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);

        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);

        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);

        const mfem::Vector densityVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(densityValue), m_layout.size(densityValue)
        );

        const mfem::Vector displacementVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(displacementValue),
            m_layout.size(displacementValue)
        );

        mfem::Vector displacementAction;

        m_preparedOperator.ApplyCompleteJacobianAction(densityVariation, displacementVariation, displacementAction);

        MFEM_VERIFY(
            displacementAction.Size() == m_layout.size(displacementResidual),
            "Prepared rotational-displacement-force MFEM adapter produced "
            "a displacement action with the wrong size."
        );

        action.SetSize(Height());
        action                   = 0.0;

        const int residualOffset = m_layout.offset(displacementResidual);

        for (int entry = 0; entry < displacementAction.Size(); ++entry) {
            action(residualOffset + entry) = displacementAction(entry);
        }
    }

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