module;

#include <algorithm>
#include <concepts>
#include <cstdint>
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include <mfem.hpp>

export module mean_field:preconditioning.gravity_field;

export import :fem;
export import :operators.context.gravity_field;
export import :preconditioning.backend_implementations;
export import :preconditioning.plan;

export namespace mean_field::preconditioning {
    struct GravityBlockDiagonal final { };
    struct GravityLowerTriangular final { };
    struct GravityUpperTriangular final { };
    struct GravityApproximateLDU final { };

    template <typename Candidate> struct IsGravityFactorizationPolicy : std::false_type { };

    template <> struct IsGravityFactorizationPolicy<GravityBlockDiagonal> : std::true_type { };
    template <> struct IsGravityFactorizationPolicy<GravityLowerTriangular> : std::true_type { };
    template <> struct IsGravityFactorizationPolicy<GravityUpperTriangular> : std::true_type { };
    template <> struct IsGravityFactorizationPolicy<GravityApproximateLDU> : std::true_type { };

    template <typename Candidate>
    concept GravityFactorizationPolicy      = IsGravityFactorizationPolicy<std::remove_cvref_t<Candidate>>::value;

    using GravityMassInverseCharacteristics = OperatorCharacteristics<
        OperatorCategory::mass_like,
        OperatorValueStructure::vector,
        OperatorSymmetry::symmetric,
        OperatorDefiniteness::positive_definite,
        OperatorRepresentation::matrix_free,
        OperatorDistribution::distributed_true_dof,
        OperatorFESpace::h_div>;

    using GravityPotentialSchurCharacteristics = OperatorCharacteristics<
        OperatorCategory::elliptic_like,
        OperatorValueStructure::scalar,
        OperatorSymmetry::symmetric,
        OperatorDefiniteness::positive_semidefinite,
        OperatorRepresentation::assembled_sparse,
        OperatorDistribution::distributed_true_dof,
        OperatorFESpace::l2,
        OperatorNullspace::constant_mode>;

    using CoupledGravityCharacteristics = OperatorCharacteristics<
        OperatorCategory::mixed,
        OperatorValueStructure::block,
        OperatorSymmetry::symmetric,
        OperatorDefiniteness::indefinite,
        OperatorRepresentation::matrix_free,
        OperatorDistribution::distributed_true_dof,
        OperatorFESpace::product>;

    namespace backend {
        template <Registered MassInverseBackend, Registered PotentialSchurBackend, GravityFactorizationPolicy Policy>
            requires Compatible<MassInverseBackend, GravityMassInverseCharacteristics> &&
                     Compatible<PotentialSchurBackend, GravityPotentialSchurCharacteristics>
        struct CoupledGravity final {
            using MassBackendType           = MassInverseBackend;
            using PotentialSchurBackendType = PotentialSchurBackend;
            using FactorizationPolicyType   = Policy;
        };

        template <Registered MassInverseBackend, Registered PotentialSchurBackend, GravityFactorizationPolicy Policy>
            requires Compatible<MassInverseBackend, GravityMassInverseCharacteristics> &&
                     Compatible<PotentialSchurBackend, GravityPotentialSchurCharacteristics>
        struct Traits<CoupledGravity<MassInverseBackend, PotentialSchurBackend, Policy>> {
            static constexpr bool registered = true;
            static constexpr ApplicationContract applicationContract =
                ::mean_field::preconditioning::backend::applicationContract<MassInverseBackend> ==
                            ApplicationContract::stationary_linear &&
                        ::mean_field::preconditioning::backend::applicationContract<PotentialSchurBackend> ==
                            ApplicationContract::stationary_linear
                    ? ApplicationContract::stationary_linear
                    : ApplicationContract::flexible;
            static constexpr bool supportsSerialExecution              = false;
            static constexpr bool supportsDistributedExecution         = true;
            static constexpr SymmetryRequirement symmetryRequirement   = SymmetryRequirement::symmetric;
            static constexpr NullspaceRequirement nullspaceRequirement = NullspaceRequirement::constant_mode_supported;
            static constexpr SurrogateRequirement surrogateRequirement = SurrogateRequirement::assembled_sparse;
            static constexpr bool requiresAssembledSparseSurrogate     = true;

            using PreparationDependencies                              = preconditioning::PreparationDependencies<
                PreparationDependency::discretization,
                PreparationDependency::geometry,
                PreparationDependency::equation_of_state,
                PreparationDependency::linearization>;

            template <OperatorCharacteristicsType Characteristics>
            static constexpr bool supports =
                Characteristics::category == OperatorCategory::mixed &&
                Characteristics::valueStructure == OperatorValueStructure::block &&
                Characteristics::symmetry == OperatorSymmetry::symmetric &&
                Characteristics::definiteness == OperatorDefiniteness::indefinite &&
                Characteristics::representation == OperatorRepresentation::matrix_free &&
                Characteristics::distribution == OperatorDistribution::distributed_true_dof &&
                Characteristics::finiteElementSpace == OperatorFESpace::product;
        };
    } // namespace backend

    template <
        backend::Registered MassBackendT,
        backend::Registered PotentialSchurBackendT,
        GravityFactorizationPolicy FactorizationPolicyT>
        requires backend::Compatible<MassBackendT, GravityMassInverseCharacteristics> &&
                 backend::Compatible<PotentialSchurBackendT, GravityPotentialSchurCharacteristics>
    class GravityFieldBlock final {
    public:
        using CorrectionBlocks =
            utils::blocks::type_list<utils::blocks::gravity::gradient::value, utils::blocks::gravity::poisson::value>;
        using ResidualBlocks = utils::blocks::
            type_list<utils::blocks::gravity::gradient::residual, utils::blocks::gravity::poisson::residual>;
        using RequiredCouplings = utils::blocks::type_list<
            Coupling<utils::blocks::gravity::gradient::residual, utils::blocks::gravity::gradient::value>,
            Coupling<utils::blocks::gravity::gradient::residual, utils::blocks::gravity::poisson::value>,
            Coupling<utils::blocks::gravity::poisson::residual, utils::blocks::gravity::gradient::value>>;
        using OperatorDescription = CoupledGravityCharacteristics;
        using BackendType         = backend::CoupledGravity<MassBackendT, PotentialSchurBackendT, FactorizationPolicyT>;
        using PreparationDependencies = typename backend::Traits<BackendType>::PreparationDependencies;
        using MassBackend             = MassBackendT;
        using PotentialSchurBackend   = PotentialSchurBackendT;
        using Factorization           = FactorizationPolicyT;

        constexpr GravityFieldBlock(
            MassBackendT massInverseBackend              = {},
            PotentialSchurBackendT potentialSchurBackend = {},
            FactorizationPolicyT factorizationPolicy     = {}
        )
            : m_massInverseBackend(std::move(massInverseBackend)),
              m_potentialSchurBackend(std::move(potentialSchurBackend)),
              m_factorizationPolicy(std::move(factorizationPolicy)) {
        }

        [[nodiscard]] constexpr const MassBackendT &massInverseBackend() const noexcept {
            return m_massInverseBackend;
        }

        [[nodiscard]] constexpr const PotentialSchurBackendT &potentialSchurBackend() const noexcept {
            return m_potentialSchurBackend;
        }

        [[nodiscard]] constexpr const FactorizationPolicyT &factorizationPolicy() const noexcept {
            return m_factorizationPolicy;
        }

    private:
        MassBackendT m_massInverseBackend;
        PotentialSchurBackendT m_potentialSchurBackend;
        FactorizationPolicyT m_factorizationPolicy;
    };

    template <
        typename MassInverseBackend,
        typename PotentialSchurBackend,
        typename FactorizationPolicy>
    GravityFieldBlock(
        MassInverseBackend,
        PotentialSchurBackend,
        FactorizationPolicy
    )
        -> GravityFieldBlock<
            MassInverseBackend,
            PotentialSchurBackend,
            FactorizationPolicy>;

    struct GravityFactorizationStatistics final {
        std::uint64_t applications{0};
        std::uint64_t massInverseApplications{0};
        std::uint64_t potentialSchurApplications{0};
        std::uint64_t divergenceApplications{0};
        std::uint64_t transposeDivergenceApplications{0};
    };

    template <GravityFactorizationPolicy Policy> class GravityFactorizationOperator final : public mfem::Solver {
    public:
        GravityFactorizationOperator(
            Policy policy,
            const mfem::Solver &massInverse,
            const mfem::Solver &potentialSchurInverse,
            const mfem::Operator &divergence
        )
            : mfem::Solver(massInverse.Height() + potentialSchurInverse.Height()),
              m_policy(std::move(policy)),
              m_massInverse(std::addressof(massInverse)),
              m_potentialSchurInverse(std::addressof(potentialSchurInverse)),
              m_divergence(std::addressof(divergence)),
              m_offsets(3),
              m_potentialWorkspace(potentialSchurInverse.Height()),
              m_gradientWorkspace(massInverse.Height()),
              m_massCorrection(massInverse.Height()) {
            if (massInverse.Height() <= 0 || massInverse.Height() != massInverse.Width()) {
                throw std::invalid_argument("The gravity factorization requires a square gradient-mass inverse.");
            }
            if (potentialSchurInverse.Height() <= 0 ||
                potentialSchurInverse.Height() != potentialSchurInverse.Width()) {
                throw std::invalid_argument("The gravity factorization requires a square potential-Schur inverse.");
            }
            if (divergence.Width() != massInverse.Width() || divergence.Height() != potentialSchurInverse.Width()) {
                throw std::invalid_argument("The gravity divergence does not connect the supplied inverse blocks.");
            }

            m_offsets[0] = 0;
            m_offsets[1] = massInverse.Height();
            m_offsets[2] = Height();
        }

        GravityFactorizationOperator(const GravityFactorizationOperator &)            = delete;
        GravityFactorizationOperator &operator=(const GravityFactorizationOperator &) = delete;
        GravityFactorizationOperator(GravityFactorizationOperator &&)                 = delete;
        GravityFactorizationOperator &operator=(GravityFactorizationOperator &&)      = delete;

        void SetOperator(const mfem::Operator &operation) override {
            if (operation.Height() != Height() || operation.Width() != Width()) {
                throw std::invalid_argument("The gravity factorization received an operator of incompatible size.");
            }
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            if (rightHandSide.Size() != Width() || action.Size() != Height()) {
                throw std::invalid_argument(
                    "The gravity factorization requires compatible, preallocated input and output vectors."
                );
            }

            const mfem::Vector gradientRightHandSide(
                const_cast<mfem::real_t *>(rightHandSide.GetData()) + m_offsets[0], m_offsets[1] - m_offsets[0]
            );
            const mfem::Vector potentialRightHandSide(
                const_cast<mfem::real_t *>(rightHandSide.GetData()) + m_offsets[1], m_offsets[2] - m_offsets[1]
            );
            mfem::Vector gradientAction(action.GetData() + m_offsets[0], m_offsets[1] - m_offsets[0]);
            mfem::Vector potentialAction(action.GetData() + m_offsets[1], m_offsets[2] - m_offsets[1]);

            if constexpr (std::same_as<Policy, GravityBlockDiagonal>) {
                m_massInverse->Mult(gradientRightHandSide, gradientAction);
                m_potentialSchurInverse->Mult(potentialRightHandSide, potentialAction);
                ++m_statistics.massInverseApplications;
                ++m_statistics.potentialSchurApplications;
            } else if constexpr (std::same_as<Policy, GravityLowerTriangular>) {
                m_massInverse->Mult(gradientRightHandSide, gradientAction);
                m_divergence->Mult(gradientAction, m_potentialWorkspace);
                m_potentialWorkspace -= potentialRightHandSide;
                m_potentialSchurInverse->Mult(m_potentialWorkspace, potentialAction);
                ++m_statistics.massInverseApplications;
                ++m_statistics.divergenceApplications;
                ++m_statistics.potentialSchurApplications;
            } else if constexpr (std::same_as<Policy, GravityUpperTriangular>) {
                m_potentialWorkspace = potentialRightHandSide;
                m_potentialWorkspace *= -1.0;
                m_potentialSchurInverse->Mult(m_potentialWorkspace, potentialAction);
                m_divergence->MultTranspose(potentialAction, m_gradientWorkspace);
                m_gradientWorkspace *= -1.0;
                m_gradientWorkspace += gradientRightHandSide;
                m_massInverse->Mult(m_gradientWorkspace, gradientAction);
                ++m_statistics.potentialSchurApplications;
                ++m_statistics.transposeDivergenceApplications;
                ++m_statistics.massInverseApplications;
            } else {
                static_assert(std::same_as<Policy, GravityApproximateLDU>);
                m_massInverse->Mult(gradientRightHandSide, gradientAction);
                m_divergence->Mult(gradientAction, m_potentialWorkspace);
                m_potentialWorkspace -= potentialRightHandSide;
                m_potentialSchurInverse->Mult(m_potentialWorkspace, potentialAction);
                m_divergence->MultTranspose(potentialAction, m_gradientWorkspace);
                m_massInverse->Mult(m_gradientWorkspace, m_massCorrection);
                gradientAction -= m_massCorrection;
                m_statistics.massInverseApplications += 2;
                ++m_statistics.divergenceApplications;
                ++m_statistics.potentialSchurApplications;
                ++m_statistics.transposeDivergenceApplications;
            }

            ++m_statistics.applications;
        }

        [[nodiscard]] const mfem::Array<int> &GetOffsets() const noexcept {
            return m_offsets;
        }

        [[nodiscard]] const GravityFactorizationStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }

    private:
        Policy m_policy;
        const mfem::Solver *m_massInverse;
        const mfem::Solver *m_potentialSchurInverse;
        const mfem::Operator *m_divergence;
        mfem::Array<int> m_offsets;
        mutable mfem::Vector m_potentialWorkspace;
        mutable mfem::Vector m_gradientWorkspace;
        mutable mfem::Vector m_massCorrection;
        mutable GravityFactorizationStatistics m_statistics;
    };

    class ReducedGravityDivergenceOperator final : public mfem::Operator {
    public:
        ReducedGravityDivergenceOperator(
            const mfem::Operator &trueDofDivergence,
            field::FieldDofMap gradientMap,
            field::FieldDofMap potentialMap
        )
            : mfem::Operator(
                  potentialMap.reduced_size(),
                  gradientMap.reduced_size()
              ),
              m_trueDofDivergence(std::addressof(trueDofDivergence)),
              m_gradientMap(std::move(gradientMap)),
              m_potentialMap(std::move(potentialMap)),
              m_gradientTrue(m_gradientMap.full_size()),
              m_potentialTrue(m_potentialMap.full_size()) {
            VerifyOperator(trueDofDivergence);
        }

        void Rebind(const mfem::Operator &trueDofDivergence) {
            VerifyOperator(trueDofDivergence);
            m_trueDofDivergence = std::addressof(trueDofDivergence);
        }

        void Mult(
            const mfem::Vector &gradient,
            mfem::Vector &potentialAction
        ) const override {
            if (gradient.Size() != Width() || potentialAction.Size() != Height()) {
                throw std::invalid_argument("The reduced gravity divergence received incompatible vectors.");
            }
            m_gradientMap.scatter(gradient, m_gradientTrue);
            m_trueDofDivergence->Mult(m_gradientTrue, m_potentialTrue);
            m_potentialMap.gather(m_potentialTrue, potentialAction);
        }

        void MultTranspose(
            const mfem::Vector &potential,
            mfem::Vector &gradientAction
        ) const override {
            if (potential.Size() != Height() || gradientAction.Size() != Width()) {
                throw std::invalid_argument("The reduced transpose divergence received incompatible vectors.");
            }
            m_potentialMap.scatter(potential, m_potentialTrue);
            m_trueDofDivergence->MultTranspose(m_potentialTrue, m_gradientTrue);
            m_gradientMap.gather(m_gradientTrue, gradientAction);
        }

    private:
        void VerifyOperator(const mfem::Operator &operation) const {
            if (operation.Width() != m_gradientMap.full_size() || operation.Height() != m_potentialMap.full_size()) {
                throw std::invalid_argument("The true-DOF divergence is incompatible with the gravity field maps.");
            }
        }

        const mfem::Operator *m_trueDofDivergence;
        field::FieldDofMap m_gradientMap;
        field::FieldDofMap m_potentialMap;
        mutable mfem::Vector m_gradientTrue;
        mutable mfem::Vector m_potentialTrue;
    };

    class ReducedFieldSolverAdapter final : public mfem::Solver {
    public:
        ReducedFieldSolverAdapter(
            const mfem::Solver &trueDofSolver,
            field::FieldDofMap map
        )
            : mfem::Solver(map.reduced_size()),
              m_trueDofSolver(std::addressof(trueDofSolver)),
              m_map(std::move(map)),
              m_rightHandSideTrue(m_map.full_size()),
              m_actionTrue(m_map.full_size()) {
            if (trueDofSolver.Height() != m_map.full_size() || trueDofSolver.Width() != m_map.full_size()) {
                throw std::invalid_argument("The true-DOF solver is incompatible with the reduced field map.");
            }
        }

        void SetOperator(const mfem::Operator &operation) override {
            if (operation.Height() != Height() || operation.Width() != Width()) {
                throw std::invalid_argument("The reduced field solver received an operator of incompatible size.");
            }
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            if (rightHandSide.Size() != Width() || action.Size() != Height()) {
                throw std::invalid_argument("The reduced field solver received incompatible vectors.");
            }
            m_map.scatter(rightHandSide, m_rightHandSideTrue);
            m_trueDofSolver->Mult(m_rightHandSideTrue, m_actionTrue);
            m_map.gather(m_actionTrue, action);
        }

    private:
        const mfem::Solver *m_trueDofSolver;
        field::FieldDofMap m_map;
        mutable mfem::Vector m_rightHandSideTrue;
        mutable mfem::Vector m_actionTrue;
    };

    [[nodiscard]] std::unique_ptr<mfem::HypreParMatrix> assembleGravityDivergenceSurrogate(const fem::FEM &f);

    [[nodiscard]] std::unique_ptr<mfem::HypreParMatrix> assembleGravityPotentialSchurSurrogate(
        const fem::FEM &f,
        const mfem::Vector &trueMassDiagonal
    );

    struct GravityFieldBlockPreparationReport final {
        bool discretizationChanged{false};
        bool geometryChanged{false};
        bool rebuiltMassInverse{false};
        bool rebuiltDivergenceBinding{false};
        bool rebuiltPotentialSchur{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return rebuiltMassInverse || rebuiltDivergenceBinding || rebuiltPotentialSchur;
        }
    };

    struct PreparedGravityFieldBlockStatistics final {
        std::uint64_t setups{0};
        std::uint64_t refreshChecks{0};
        std::uint64_t refreshes{0};
        std::uint64_t noOpRefreshes{0};
    };

    template <typename Candidate>
    concept ImplementedGravityMassBackend = std::same_as<std::remove_cvref_t<Candidate>, backend::Diagonal> ||
                                            std::same_as<std::remove_cvref_t<Candidate>, backend::MatrixFreeChebyshev>;

    template <backend::Registered MassBackend, backend::ApplicationMode Mode, GravityFactorizationPolicy Policy>
        requires ImplementedGravityMassBackend<MassBackend> &&
                 backend::Compatible<MassBackend, GravityMassInverseCharacteristics>
    class PreparedGravityFieldBlock final : public mfem::Solver {
    public:
        using Block               = GravityFieldBlock<MassBackend, backend::HypreBoomerAMG<Mode>, Policy>;
        using PreparedMassInverse = std::conditional_t<
            std::same_as<MassBackend, backend::Diagonal>,
            backend::PreparedDiagonal,
            backend::PreparedMatrixFreeChebyshev>;

        PreparedGravityFieldBlock(
            const fem::FEM &f,
            const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext,
            Block block
        )
            : mfem::Solver(GravitySize(geometryContext)),
              m_block(std::move(block)),
              m_geometryContext(std::addressof(geometryContext)),
              m_gradientMap(geometryContext.GetMassOperator().GetFluxMap()),
              m_potentialMap(geometryContext.GetSourceOperator().GetPotentialMap()),
              m_divergence(
                  geometryContext.GetDivergenceOperator(),
                  m_gradientMap,
                  m_potentialMap
              ),
              m_massInverse(MakeMassInverse(
                  f,
                  geometryContext,
                  m_block.massInverseBackend()
              )),
              m_potentialSchurSurrogate(AssemblePotentialSchur(
                  f,
                  geometryContext
              )),
              m_potentialSchurInverse(
                  m_block.potentialSchurBackend(),
                  *m_potentialSchurSurrogate
              ),
              m_reducedPotentialSchurInverse(
                  m_potentialSchurInverse,
                  m_potentialMap
              ),
              m_factorization(
                  m_block.factorizationPolicy(),
                  m_massInverse,
                  m_reducedPotentialSchurInverse,
                  m_divergence
              ),
              m_discretizationRevision(geometryContext.GetDiscretizationRevision()),
              m_displacementRevision(geometryContext.GetDisplacementRevision()) {
            if (!geometryContext.IsPrepared()) {
                throw std::logic_error("The gravity field block requires a prepared gravity geometry context.");
            }
            m_statistics.setups = 1;
        }

        PreparedGravityFieldBlock(const PreparedGravityFieldBlock &)            = delete;
        PreparedGravityFieldBlock &operator=(const PreparedGravityFieldBlock &) = delete;
        PreparedGravityFieldBlock(PreparedGravityFieldBlock &&)                 = delete;
        PreparedGravityFieldBlock &operator=(PreparedGravityFieldBlock &&)      = delete;

        void SetOperator(const mfem::Operator &operation) override {
            m_factorization.SetOperator(operation);
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            if (!IsCurrent()) {
                throw std::logic_error("The gravity field block is stale; refresh it before application.");
            }
            m_factorization.Mult(rightHandSide, action);
        }

        [[nodiscard]] bool IsCurrent() const noexcept {
            return m_geometryContext->IsPrepared() &&
                   m_geometryContext->GetDiscretizationRevision() == m_discretizationRevision &&
                   m_geometryContext->GetDisplacementRevision() == m_displacementRevision;
        }

        [[nodiscard]] GravityFieldBlockPreparationReport Refresh(
            const fem::FEM &f,
            const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext
        ) {
            if (!geometryContext.IsPrepared()) {
                throw std::logic_error("The gravity field block cannot refresh from unprepared geometry.");
            }
            if (std::addressof(geometryContext) != m_geometryContext) {
                throw std::invalid_argument("A prepared gravity field block cannot change geometry-context identity.");
            }

            ++m_statistics.refreshChecks;
            GravityFieldBlockPreparationReport report{
                .discretizationChanged = geometryContext.GetDiscretizationRevision() != m_discretizationRevision,
                .geometryChanged       = geometryContext.GetDisplacementRevision() != m_displacementRevision
            };
            if (!report.discretizationChanged && !report.geometryChanged) {
                ++m_statistics.noOpRefreshes;
                return report;
            }

            m_divergence.Rebind(geometryContext.GetDivergenceOperator());
            report.rebuiltDivergenceBinding = report.discretizationChanged;

            RefreshMassInverse(geometryContext);
            report.rebuiltMassInverse = true;

            auto potentialSchur       = AssemblePotentialSchur(f, geometryContext);
            m_potentialSchurInverse.Refresh(*potentialSchur);
            m_potentialSchurSurrogate    = std::move(potentialSchur);
            report.rebuiltPotentialSchur = true;

            m_discretizationRevision     = geometryContext.GetDiscretizationRevision();
            m_displacementRevision       = geometryContext.GetDisplacementRevision();
            ++m_statistics.refreshes;
            return report;
        }

        [[nodiscard]] const Block &GetBlock() const noexcept {
            return m_block;
        }

        [[nodiscard]] const mfem::Array<int> &GetOffsets() const noexcept {
            return m_factorization.GetOffsets();
        }

        [[nodiscard]] const PreparedMassInverse &GetMassInverse() const {
            if (!IsCurrent()) {
                throw std::logic_error("The gravity mass inverse is stale; refresh its owning gravity block first.");
            }
            return m_massInverse;
        }

        [[nodiscard]] const backend::PreparedHypreBoomerAMG<Mode> &GetPotentialSchurInverse() const noexcept {
            return m_potentialSchurInverse;
        }

        [[nodiscard]] const mfem::HypreParMatrix &GetPotentialSchurSurrogate() const noexcept {
            return *m_potentialSchurSurrogate;
        }

        [[nodiscard]] const GravityFactorizationOperator<Policy> &GetFactorization() const noexcept {
            return m_factorization;
        }

        [[nodiscard]] const PreparedGravityFieldBlockStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }

    private:
        [[nodiscard]] static int
        GravitySize(const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext) {
            if (!geometryContext.IsPrepared()) {
                throw std::logic_error("The gravity field block requires a prepared gravity geometry context.");
            }
            return geometryContext.GetMassOperator().GetFluxMap().reduced_size() +
                   geometryContext.GetSourceOperator().GetPotentialMap().reduced_size();
        }

        [[nodiscard]] static mfem::Vector AssembleReducedMassDiagonal(
            const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext
        ) {
            mfem::Vector diagonal;
            geometryContext.GetMassOperator().AssembleDiagonal(diagonal);
            return diagonal;
        }

        [[nodiscard]] static PreparedMassInverse MakeMassInverse(
            const fem::FEM &f,
            const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext,
            const MassBackend &backendConfiguration
        ) {
            if constexpr (std::same_as<MassBackend, backend::Diagonal>) {
                return PreparedMassInverse{backendConfiguration, AssembleReducedMassDiagonal(geometryContext)};
            } else {
                static_assert(std::same_as<MassBackend, backend::MatrixFreeChebyshev>);
                return PreparedMassInverse{
                    backendConfiguration, geometryContext.GetMassOperator(), f.gravityFluxFes->GetComm()
                };
            }
        }

        void RefreshMassInverse(const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext) {
            if constexpr (std::same_as<MassBackend, backend::Diagonal>) {
                m_massInverse.Refresh(AssembleReducedMassDiagonal(geometryContext));
            } else {
                static_assert(std::same_as<MassBackend, backend::MatrixFreeChebyshev>);
                m_massInverse.Refresh(geometryContext.GetMassOperator());
            }
        }

        [[nodiscard]] static std::unique_ptr<mfem::HypreParMatrix> AssemblePotentialSchur(
            const fem::FEM &f,
            const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext
        ) {
            mfem::Vector trueMassDiagonal;
            geometryContext.GetMassOperator().AssembleTrueDiagonal(trueMassDiagonal);
            return assembleGravityPotentialSchurSurrogate(f, trueMassDiagonal);
        }

        Block m_block;
        const operators::context::gravity_field::GravityFieldGeometryContext *m_geometryContext;
        field::FieldDofMap m_gradientMap;
        field::FieldDofMap m_potentialMap;
        ReducedGravityDivergenceOperator m_divergence;
        PreparedMassInverse m_massInverse;
        std::unique_ptr<mfem::HypreParMatrix> m_potentialSchurSurrogate;
        backend::PreparedHypreBoomerAMG<Mode> m_potentialSchurInverse;
        ReducedFieldSolverAdapter m_reducedPotentialSchurInverse;
        GravityFactorizationOperator<Policy> m_factorization;
        operators::context::gravity_field::DiscretizationRevision m_discretizationRevision;
        operators::context::gravity_field::DisplacementRevision m_displacementRevision;
        PreparedGravityFieldBlockStatistics m_statistics;
    };

    template <
        backend::Registered MassBackend,
        backend::ApplicationMode Mode,
        GravityFactorizationPolicy Policy>
        requires ImplementedGravityMassBackend<MassBackend> && backend::Compatible<
                                                                   MassBackend,
                                                                   GravityMassInverseCharacteristics>
    [[nodiscard]] auto prepare(
        const fem::FEM &f,
        const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext,
        GravityFieldBlock<
            MassBackend,
            backend::HypreBoomerAMG<Mode>,
            Policy> block
    ) {
        return PreparedGravityFieldBlock<MassBackend, Mode, Policy>{f, geometryContext, std::move(block)};
    }
} // namespace mean_field::preconditioning