module;

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

#include <mfem.hpp>

export module mean_field:preconditioning.stellar_structure;

export import :preconditioning.gravity_field;
export import :preconditioning.material_surface;
export import :preconditioning.stellar_equilibrium;

export namespace mean_field::preconditioning {
    struct IndependentStellarSubsystems final { };
    struct MaterialThenGravityTriangular final { };
    struct GravityThenMaterialTriangular final { };
    struct ApproximateStellarBlockLDU final { };

    template <typename Candidate> struct IsStellarStructureFactorizationPolicy : std::false_type { };
    template <> struct IsStellarStructureFactorizationPolicy<IndependentStellarSubsystems> : std::true_type { };
    template <> struct IsStellarStructureFactorizationPolicy<MaterialThenGravityTriangular> : std::true_type { };
    template <> struct IsStellarStructureFactorizationPolicy<GravityThenMaterialTriangular> : std::true_type { };
    template <> struct IsStellarStructureFactorizationPolicy<ApproximateStellarBlockLDU> : std::true_type { };

    template <typename Candidate>
    concept StellarStructureFactorizationPolicy =
        IsStellarStructureFactorizationPolicy<std::remove_cvref_t<Candidate>>::value;

    namespace detail {
        template <typename... Lists> struct StellarStructureConcatenate;

        template <> struct StellarStructureConcatenate<> {
            using Type = utils::blocks::type_list<>;
        };

        template <typename... Types> struct StellarStructureConcatenate<utils::blocks::type_list<Types...>> {
            using Type = utils::blocks::type_list<Types...>;
        };

        template <typename... Left, typename... Right, typename... Remaining>
        struct StellarStructureConcatenate<
            utils::blocks::type_list<Left...>,
            utils::blocks::type_list<Right...>,
            Remaining...> {
            using Type =
                typename StellarStructureConcatenate<utils::blocks::type_list<Left..., Right...>, Remaining...>::Type;
        };

        template <typename... Lists>
        using StellarStructureConcatenateT = typename StellarStructureConcatenate<Lists...>::Type;

        template <typename Residual, typename Corrections, typename JacobianForm>
        struct StellarStructureCouplingsForResidual;

        template <typename Residual, typename JacobianForm>
        struct StellarStructureCouplingsForResidual<Residual, utils::blocks::type_list<>, JacobianForm> {
            using Type = utils::blocks::type_list<>;
        };

        template <typename Residual, typename First, typename... Remaining, typename JacobianForm>
        struct StellarStructureCouplingsForResidual<
            Residual,
            utils::blocks::type_list<First, Remaining...>,
            JacobianForm> {
        private:
            using Tail = typename StellarStructureCouplingsForResidual<
                Residual,
                utils::blocks::type_list<Remaining...>,
                JacobianForm>::Type;

        public:
            using Type = std::conditional_t<
                utils::blocks::has_jacobian_coupling_v<Residual, First, JacobianForm>,
                StellarStructureConcatenateT<utils::blocks::type_list<Coupling<Residual, First>>, Tail>,
                Tail>;
        };

        template <typename Residuals, typename Corrections, typename JacobianForm>
        struct StellarStructureInducedCouplings;

        template <typename Corrections, typename JacobianForm>
        struct StellarStructureInducedCouplings<utils::blocks::type_list<>, Corrections, JacobianForm> {
            using Type = utils::blocks::type_list<>;
        };

        template <typename First, typename... Remaining, typename Corrections, typename JacobianForm>
        struct StellarStructureInducedCouplings<
            utils::blocks::type_list<First, Remaining...>,
            Corrections,
            JacobianForm> {
            using Type = StellarStructureConcatenateT<
                typename StellarStructureCouplingsForResidual<First, Corrections, JacobianForm>::Type,
                typename StellarStructureInducedCouplings<
                    utils::blocks::type_list<Remaining...>,
                    Corrections,
                    JacobianForm>::Type>;
        };

        template <typename Left, typename Right> struct StellarStructureListsAreDisjoint;

        template <typename... Left, typename Right>
        struct StellarStructureListsAreDisjoint<utils::blocks::type_list<Left...>, Right>
            : std::bool_constant<(!utils::blocks::contains_type_v<Left, Right> && ...)> { };

        template <typename Candidate, typename Universe> struct StellarStructureListIsSubset;

        template <typename... Candidates, typename Universe>
        struct StellarStructureListIsSubset<utils::blocks::type_list<Candidates...>, Universe>
            : std::bool_constant<(utils::blocks::contains_type_v<Candidates, Universe> && ...)> { };
    } // namespace detail

    using CoupledStellarStructureCharacteristics = OperatorCharacteristics<
        OperatorCategory::mixed,
        OperatorValueStructure::block,
        OperatorSymmetry::nonsymmetric,
        OperatorDefiniteness::unspecified,
        OperatorRepresentation::matrix_free,
        OperatorDistribution::distributed_true_dof,
        OperatorFESpace::product>;

    namespace backend {
        template <
            Registered MaterialSurfaceBackend,
            Registered GravityBackend,
            StellarStructureFactorizationPolicy Policy>
        struct CoupledStellarStructure final {
            using MaterialSurfaceBackendType = MaterialSurfaceBackend;
            using GravityBackendType         = GravityBackend;
            using FactorizationPolicyType    = Policy;
        };

        template <
            Registered MaterialSurfaceBackend,
            Registered GravityBackend,
            StellarStructureFactorizationPolicy Policy>
        struct Traits<CoupledStellarStructure<MaterialSurfaceBackend, GravityBackend, Policy>> {
            static constexpr bool registered = true;
            static constexpr ApplicationContract applicationContract =
                ::mean_field::preconditioning::backend::applicationContract<MaterialSurfaceBackend> ==
                            ApplicationContract::stationary_linear &&
                        ::mean_field::preconditioning::backend::applicationContract<GravityBackend> ==
                            ApplicationContract::stationary_linear
                    ? ApplicationContract::stationary_linear
                    : ApplicationContract::flexible;
            static constexpr bool supportsSerialExecution = Traits<MaterialSurfaceBackend>::supportsSerialExecution &&
                                                            Traits<GravityBackend>::supportsSerialExecution;
            static constexpr bool supportsDistributedExecution =
                Traits<MaterialSurfaceBackend>::supportsDistributedExecution &&
                Traits<GravityBackend>::supportsDistributedExecution;
            static constexpr SymmetryRequirement symmetryRequirement   = SymmetryRequirement::none;
            static constexpr NullspaceRequirement nullspaceRequirement = NullspaceRequirement::constant_mode_supported;
            static constexpr SurrogateRequirement surrogateRequirement = SurrogateRequirement::assembled_sparse;
            static constexpr bool requiresAssembledSparseSurrogate =
                Traits<MaterialSurfaceBackend>::requiresAssembledSparseSurrogate ||
                Traits<GravityBackend>::requiresAssembledSparseSurrogate;

            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::nonsymmetric &&
                Characteristics::representation == OperatorRepresentation::matrix_free &&
                Characteristics::distribution == OperatorDistribution::distributed_true_dof &&
                Characteristics::finiteElementSpace == OperatorFESpace::product;
        };
    } // namespace backend

    template <
        PreconditionerComponent MaterialSurfaceComponentT,
        PreconditionerComponent GravityComponentT,
        typename FormT,
        typename JacobianFormT,
        StellarStructureFactorizationPolicy PolicyT>
        requires utils::blocks::valid_jacobian_form<FormT, JacobianFormT> &&
                 detail::StellarStructureListsAreDisjoint<
                     typename MaterialSurfaceComponentT::CorrectionBlocks,
                     typename GravityComponentT::CorrectionBlocks>::value &&
                 detail::StellarStructureListsAreDisjoint<
                     typename MaterialSurfaceComponentT::ResidualBlocks,
                     typename GravityComponentT::ResidualBlocks>::value &&
                 detail::StellarStructureListIsSubset<
                     typename MaterialSurfaceComponentT::CorrectionBlocks,
                     typename FormT::value_blocks>::value &&
                 detail::StellarStructureListIsSubset<
                     typename GravityComponentT::CorrectionBlocks,
                     typename FormT::value_blocks>::value &&
                 detail::StellarStructureListIsSubset<
                     typename MaterialSurfaceComponentT::ResidualBlocks,
                     typename FormT::residual_blocks>::value &&
                 detail::StellarStructureListIsSubset<
                     typename GravityComponentT::ResidualBlocks,
                     typename FormT::residual_blocks>::value
    class StellarStructureBlock final {
    public:
        using MaterialSurfaceComponent = MaterialSurfaceComponentT;
        using GravityComponent         = GravityComponentT;
        using Form                     = FormT;
        using JacobianForm             = JacobianFormT;
        using Factorization            = PolicyT;
        using CorrectionBlocks         = detail::StellarStructureConcatenateT<
            typename MaterialSurfaceComponent::CorrectionBlocks,
            typename GravityComponent::CorrectionBlocks>;
        using ResidualBlocks = detail::StellarStructureConcatenateT<
            typename MaterialSurfaceComponent::ResidualBlocks,
            typename GravityComponent::ResidualBlocks>;
        using MaterialToGravityCouplings = typename detail::StellarStructureInducedCouplings<
            typename GravityComponent::ResidualBlocks,
            typename MaterialSurfaceComponent::CorrectionBlocks,
            JacobianForm>::Type;
        using GravityToMaterialCouplings = typename detail::StellarStructureInducedCouplings<
            typename MaterialSurfaceComponent::ResidualBlocks,
            typename GravityComponent::CorrectionBlocks,
            JacobianForm>::Type;
        using RequiredCouplings = detail::StellarStructureConcatenateT<
            typename MaterialSurfaceComponent::RequiredCouplings,
            typename GravityComponent::RequiredCouplings,
            MaterialToGravityCouplings,
            GravityToMaterialCouplings>;
        using OperatorDescription = CoupledStellarStructureCharacteristics;
        using BackendType         = backend::CoupledStellarStructure<
            typename MaterialSurfaceComponent::BackendType,
            typename GravityComponent::BackendType,
            Factorization>;
        using PreparationDependencies = typename backend::Traits<BackendType>::PreparationDependencies;

        constexpr StellarStructureBlock(
            MaterialSurfaceComponent materialSurfaceComponent,
            GravityComponent gravityComponent,
            Factorization factorization = {}
        )
            : m_materialSurfaceComponent(std::move(materialSurfaceComponent)),
              m_gravityComponent(std::move(gravityComponent)),
              m_factorization(std::move(factorization)) {
        }

        [[nodiscard]] constexpr const MaterialSurfaceComponent &materialSurfaceComponent() const noexcept {
            return m_materialSurfaceComponent;
        }

        [[nodiscard]] constexpr const GravityComponent &gravityComponent() const noexcept {
            return m_gravityComponent;
        }

        [[nodiscard]] constexpr const Factorization &factorizationPolicy() const noexcept {
            return m_factorization;
        }

    private:
        MaterialSurfaceComponent m_materialSurfaceComponent;
        GravityComponent m_gravityComponent;
        Factorization m_factorization;
    };

    template <typename Candidate>
    concept StellarStructureCrossCouplingOperator = requires(
        const Candidate &couplings,
        const mfem::Vector &materialDirection,
        const mfem::Vector &gravityDirection,
        mfem::Vector &materialAction,
        mfem::Vector &gravityAction
    ) {
        { couplings.MaterialSize() } -> std::same_as<int>;
        { couplings.GravitySize() } -> std::same_as<int>;
        couplings.ApplyMaterialToGravity(materialDirection, gravityAction);
        couplings.ApplyGravityToMaterial(gravityDirection, materialAction);
    };

    struct StellarStructureFactorizationStatistics final {
        std::uint64_t applications{0};
        std::uint64_t materialSurfaceInverseApplications{0};
        std::uint64_t gravityInverseApplications{0};
        std::uint64_t materialToGravityApplications{0};
        std::uint64_t gravityToMaterialApplications{0};
    };

    template <StellarStructureFactorizationPolicy Policy, StellarStructureCrossCouplingOperator CouplingOperator>
    class StellarStructureFactorizationOperator final : public mfem::Solver {
    public:
        StellarStructureFactorizationOperator(
            Policy policy,
            const mfem::Solver &materialSurfaceInverse,
            const mfem::Solver &gravityInverse,
            const CouplingOperator &couplings
        )
            : mfem::Solver(materialSurfaceInverse.Height() + gravityInverse.Height()),
              m_policy(std::move(policy)),
              m_materialSurfaceInverse(std::addressof(materialSurfaceInverse)),
              m_gravityInverse(std::addressof(gravityInverse)),
              m_couplings(std::addressof(couplings)),
              m_materialWorkspace(materialSurfaceInverse.Height()),
              m_gravityWorkspace(gravityInverse.Height()) {
            if (materialSurfaceInverse.Height() <= 0 ||
                materialSurfaceInverse.Height() != materialSurfaceInverse.Width() || gravityInverse.Height() <= 0 ||
                gravityInverse.Height() != gravityInverse.Width() ||
                materialSurfaceInverse.Height() != couplings.MaterialSize() ||
                gravityInverse.Height() != couplings.GravitySize()) {
                throw std::invalid_argument(
                    "The stellar-structure inverse blocks do not match the cross-coupling operator."
                );
            }
        }

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

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

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

            action = 0.0;
            const mfem::Vector materialRightHandSide(
                const_cast<mfem::real_t *>(rightHandSide.GetData()), m_materialSurfaceInverse->Width()
            );
            const mfem::Vector gravityRightHandSide(
                const_cast<mfem::real_t *>(rightHandSide.GetData()) + m_materialSurfaceInverse->Width(),
                m_gravityInverse->Width()
            );
            mfem::Vector materialAction(action, 0, m_materialSurfaceInverse->Height());
            mfem::Vector gravityAction(action, m_materialSurfaceInverse->Height(), m_gravityInverse->Height());

            if constexpr (std::same_as<Policy, IndependentStellarSubsystems>) {
                m_materialSurfaceInverse->Mult(materialRightHandSide, materialAction);
                m_gravityInverse->Mult(gravityRightHandSide, gravityAction);
                ++m_statistics.materialSurfaceInverseApplications;
                ++m_statistics.gravityInverseApplications;
            } else if constexpr (std::same_as<Policy, MaterialThenGravityTriangular>) {
                m_materialSurfaceInverse->Mult(materialRightHandSide, materialAction);
                m_couplings->ApplyMaterialToGravity(materialAction, m_gravityWorkspace);
                m_gravityWorkspace *= -1.0;
                m_gravityWorkspace += gravityRightHandSide;
                m_gravityInverse->Mult(m_gravityWorkspace, gravityAction);
                ++m_statistics.materialSurfaceInverseApplications;
                ++m_statistics.materialToGravityApplications;
                ++m_statistics.gravityInverseApplications;
            } else if constexpr (std::same_as<Policy, GravityThenMaterialTriangular>) {
                m_gravityInverse->Mult(gravityRightHandSide, gravityAction);
                m_couplings->ApplyGravityToMaterial(gravityAction, m_materialWorkspace);
                m_materialWorkspace *= -1.0;
                m_materialWorkspace += materialRightHandSide;
                m_materialSurfaceInverse->Mult(m_materialWorkspace, materialAction);
                ++m_statistics.gravityInverseApplications;
                ++m_statistics.gravityToMaterialApplications;
                ++m_statistics.materialSurfaceInverseApplications;
            } else {
                static_assert(std::same_as<Policy, ApproximateStellarBlockLDU>);
                m_materialSurfaceInverse->Mult(materialRightHandSide, materialAction);
                m_couplings->ApplyMaterialToGravity(materialAction, m_gravityWorkspace);
                m_gravityWorkspace *= -1.0;
                m_gravityWorkspace += gravityRightHandSide;
                m_gravityInverse->Mult(m_gravityWorkspace, gravityAction);
                m_couplings->ApplyGravityToMaterial(gravityAction, m_materialWorkspace);
                m_materialWorkspace *= -1.0;
                m_materialWorkspace += materialRightHandSide;
                m_materialSurfaceInverse->Mult(m_materialWorkspace, materialAction);
                m_statistics.materialSurfaceInverseApplications += 2;
                ++m_statistics.materialToGravityApplications;
                ++m_statistics.gravityInverseApplications;
                ++m_statistics.gravityToMaterialApplications;
            }
            materialAction.SyncAliasMemory(action);
            gravityAction.SyncAliasMemory(action);
            ++m_statistics.applications;
        }

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

    private:
        Policy m_policy;
        const mfem::Solver *m_materialSurfaceInverse;
        const mfem::Solver *m_gravityInverse;
        const CouplingOperator *m_couplings;
        mutable mfem::Vector m_materialWorkspace;
        mutable mfem::Vector m_gravityWorkspace;
        mutable StellarStructureFactorizationStatistics m_statistics;
    };

    class StellarStructureCrossJacobianOperator final : public mfem::Operator {
    public:
        explicit StellarStructureCrossJacobianOperator(const operators::PreparedStellarEquilibriumOperator &operation)
            : mfem::Operator(MaterialSizeOf(operation) + GravitySizeOf(operation)),
              m_operation(std::addressof(operation)),
              m_materialOffsets(4),
              m_gravityOffsets(3),
              m_combinedOffsets(3),
              m_gravityDirection(operation.GetGravityJacobianOperator().Width()),
              m_volumeDisplacement(operation.GetDomainDeformation().volumeDisplacementSize()),
              m_mechanicalAction(operation.GetDomainDeformation().volumeDisplacementSize()),
              m_zeroEnthalpy(operation.GetBarotropicClosureOperator().GetEnthalpySize()) {
            const auto &context  = operation.GetGravityContext();
            m_materialOffsets[0] = 0;
            m_materialOffsets[1] = context.GetDensityMap().reduced_size();
            m_materialOffsets[2] = m_materialOffsets[1] + operation.GetDomainDeformation().parameterCount();
            m_materialOffsets[3] = MaterialSizeOf(operation);
            m_gravityOffsets[0]  = 0;
            m_gravityOffsets[1]  = context.GetGravityGradientMap().reduced_size();
            m_gravityOffsets[2]  = GravitySizeOf(operation);
            m_combinedOffsets[0] = 0;
            m_combinedOffsets[1] = MaterialSize();
            m_combinedOffsets[2] = Height();
            m_zeroEnthalpy       = 0.0;
        }

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override {
            VerifyCombined(direction, action);
            action = 0.0;
            const mfem::Vector materialDirection(const_cast<mfem::real_t *>(direction.GetData()), MaterialSize());
            const mfem::Vector gravityDirection(
                const_cast<mfem::real_t *>(direction.GetData()) + MaterialSize(), GravitySize()
            );
            mfem::Vector materialAction(action, 0, MaterialSize());
            mfem::Vector gravityAction(action, MaterialSize(), GravitySize());
            ApplyMaterialToGravity(materialDirection, gravityAction);
            ApplyGravityToMaterial(gravityDirection, materialAction);
            materialAction.SyncAliasMemory(action);
            gravityAction.SyncAliasMemory(action);
        }

        void ApplyMaterialToGravity(
            const mfem::Vector &materialDirection,
            mfem::Vector &gravityAction
        ) const {
            VerifyMaterial(materialDirection, "direction");
            VerifyGravity(gravityAction, "action");
            const auto densityDirection = MaterialBlock(materialDirection, 0);
            const auto surfaceDirection = MaterialBlock(materialDirection, 1);
            const auto &gravityOffsets  = m_operation->GetGravityOperator().GetStateOffsets();
            using GravityForm           = utils::blocks::gravity_field_form;
            constexpr auto densityBlock =
                utils::blocks::get_value_block<GravityForm>(utils::blocks::density_field.mass_term);
            constexpr auto displacementBlock =
                utils::blocks::get_value_block<GravityForm>(utils::blocks::displacement_field.geometry_term);

            m_gravityDirection          = 0.0;
            auto packedDensityDirection = MutableBlock(m_gravityDirection, gravityOffsets, densityBlock.index);
            packedDensityDirection      = densityDirection;
            packedDensityDirection.SyncAliasMemory(m_gravityDirection);
            m_operation->GetDomainDeformation().applyJacobian(
                m_operation->GetSurfaceDeformationParameters(), surfaceDirection, m_volumeDisplacement
            );
            auto packedDisplacementDirection =
                MutableBlock(m_gravityDirection, gravityOffsets, displacementBlock.index);
            packedDisplacementDirection = m_volumeDisplacement;
            packedDisplacementDirection.SyncAliasMemory(m_gravityDirection);
            m_operation->GetGravityJacobianOperator().Mult(m_gravityDirection, gravityAction);
        }

        void ApplyGravityToMaterial(
            const mfem::Vector &gravityDirection,
            mfem::Vector &materialAction
        ) const {
            VerifyGravity(gravityDirection, "direction");
            VerifyMaterial(materialAction, "action");
            const auto gravityGradientDirection  = GravityBlock(gravityDirection, 0);
            const auto gravityPotentialDirection = GravityBlock(gravityDirection, 1);
            auto densityAction                   = MaterialBlock(materialAction, 0);
            auto surfaceAction                   = MaterialBlock(materialAction, 1);
            auto enthalpyAction                  = MaterialBlock(materialAction, 2);

            densityAction                        = 0.0;
            m_operation->GetDisplacementOperator().ApplyGravityGradientJacobianAction(
                gravityGradientDirection, m_mechanicalAction
            );
            m_operation->GetDomainDeformation().applyJacobianTranspose(
                m_operation->GetSurfaceDeformationParameters(), m_mechanicalAction, surfaceAction
            );
            m_operation->GetHydrostaticOperator().ApplyGravityPotentialJacobianAction(
                gravityPotentialDirection, enthalpyAction
            );
            m_operation->GetSurfaceConstraintOperator().ApplyJacobianRows(m_zeroEnthalpy, enthalpyAction);
            densityAction.SyncAliasMemory(materialAction);
            surfaceAction.SyncAliasMemory(materialAction);
            enthalpyAction.SyncAliasMemory(materialAction);
        }

        [[nodiscard]] int MaterialSize() const noexcept {
            return m_materialOffsets.Last();
        }

        [[nodiscard]] int GravitySize() const noexcept {
            return m_gravityOffsets.Last();
        }

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

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

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

    private:
        [[nodiscard]] static int MaterialSizeOf(const operators::PreparedStellarEquilibriumOperator &operation) {
            if (!operation.IsPrepared()) {
                throw std::logic_error("The stellar-structure cross Jacobian requires a prepared operator.");
            }
            return operation.GetGravityContext().GetDensityMap().reduced_size() +
                   operation.GetDomainDeformation().parameterCount() +
                   operation.GetBarotropicClosureOperator().GetEnthalpySize();
        }

        [[nodiscard]] static int GravitySizeOf(const operators::PreparedStellarEquilibriumOperator &operation) {
            return operation.GetGravityContext().GetGravityGradientMap().reduced_size() +
                   operation.GetGravityContext().GetGravityPotentialMap().reduced_size();
        }

        [[nodiscard]] static mfem::Vector MutableBlock(
            mfem::Vector &vector,
            const mfem::Array<int> &offsets,
            const int block
        ) {
            return mfem::Vector(vector, offsets[block], offsets[block + 1] - offsets[block]);
        }

        [[nodiscard]] mfem::Vector MaterialBlock(
            const mfem::Vector &vector,
            const int block
        ) const {
            return mfem::Vector(
                const_cast<mfem::real_t *>(vector.GetData()) + m_materialOffsets[block],
                m_materialOffsets[block + 1] - m_materialOffsets[block]
            );
        }

        [[nodiscard]] mfem::Vector MaterialBlock(
            mfem::Vector &vector,
            const int block
        ) const {
            return mfem::Vector(
                vector, m_materialOffsets[block], m_materialOffsets[block + 1] - m_materialOffsets[block]
            );
        }

        [[nodiscard]] mfem::Vector GravityBlock(
            const mfem::Vector &vector,
            const int block
        ) const {
            return mfem::Vector(
                const_cast<mfem::real_t *>(vector.GetData()) + m_gravityOffsets[block],
                m_gravityOffsets[block + 1] - m_gravityOffsets[block]
            );
        }

        void VerifyCombined(
            const mfem::Vector &direction,
            const mfem::Vector &action
        ) const {
            if (direction.Size() != Width() || action.Size() != Height()) {
                throw std::invalid_argument(
                    "The stellar-structure cross Jacobian requires compatible, preallocated vectors."
                );
            }
        }

        void VerifyMaterial(
            const mfem::Vector &vector,
            const char *role
        ) const {
            if (vector.Size() != MaterialSize()) {
                throw std::invalid_argument(
                    std::string("The stellar-structure material ") + role + " has the wrong size."
                );
            }
        }

        void VerifyGravity(
            const mfem::Vector &vector,
            const char *role
        ) const {
            if (vector.Size() != GravitySize()) {
                throw std::invalid_argument(
                    std::string("The stellar-structure gravity ") + role + " has the wrong size."
                );
            }
        }

        const operators::PreparedStellarEquilibriumOperator *m_operation;
        mfem::Array<int> m_materialOffsets;
        mfem::Array<int> m_gravityOffsets;
        mfem::Array<int> m_combinedOffsets;
        mutable mfem::Vector m_gravityDirection;
        mutable mfem::Vector m_volumeDisplacement;
        mutable mfem::Vector m_mechanicalAction;
        mfem::Vector m_zeroEnthalpy;
    };

    struct StellarStructureBlockPreparationReport final {
        MaterialSurfaceBlockPreparationReport materialSurface;
        GravityFieldBlockPreparationReport gravity;

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return materialSurface.DidAnyWork() || gravity.DidAnyWork();
        }
    };

    /*
     * Material/surface execution and stellar-structure execution are separate
     * capabilities.  The latter also owns the physical cross-Jacobian and
     * gravity-context wiring, which currently target the legacy prepared core.
     * Add future cores here only together with matching cross-coupling and
     * preparation implementations.
     */
    template <typename PhysicalCore> struct StellarStructureExecutableRuntime {
        static constexpr bool available = false;
    };

    template <> struct StellarStructureExecutableRuntime<operators::PreparedStellarEquilibriumOperator> {
        static constexpr bool available = true;
    };

    template <typename PhysicalCore>
    concept ExecutableStellarStructureRuntimeFor = requires {
        {
            StellarStructureExecutableRuntime<std::remove_cvref_t<PhysicalCore>>::available
        } -> std::convertible_to<bool>;
        requires StellarStructureExecutableRuntime<std::remove_cvref_t<PhysicalCore>>::available;
    };

    template <typename Descriptor, typename PhysicalCore>
    concept StellarStructureRuntimeFor =
        MaterialSurfaceRuntimeFor<Descriptor, PhysicalCore> && ExecutableStellarStructureRuntimeFor<PhysicalCore>;

    template <typename Candidate>
    concept StellarStructurePreconditionerProblem =
        equilibrium::DiscretizedStellarEquilibriumProblem<Candidate> &&
        DefaultStellarStructurePhysicalTopologySupportedFor<typename std::remove_cvref_t<Candidate>::ModelType> &&
        requires {
            requires StellarStructureRuntimeFor<
                MaterialSurfaceDescriptorFor<std::remove_cvref_t<Candidate>>,
                typename std::remove_cvref_t<Candidate>::PhysicalCoreType>;
        };

    namespace detail {
        template <StellarStructurePreconditionerProblem Problem>
        [[nodiscard]] const auto &physicalOperator(const Problem &problem) {
            return problem.GetPhysicalOperator();
        }
    } // namespace detail

    template <typename Problem, typename MaterialComponent, typename GravityComponent>
    concept StellarStructurePreparableFor =
        StellarStructurePreconditionerProblem<std::remove_cvref_t<Problem>> &&
        PreconditionerComponent<std::remove_cvref_t<MaterialComponent>> &&
        PreconditionerComponent<std::remove_cvref_t<GravityComponent>> &&
        requires(
            const std::remove_cvref_t<Problem> &problem,
            std::remove_cvref_t<MaterialComponent> materialComponent,
            std::remove_cvref_t<GravityComponent> gravityComponent
        ) {
            preconditioning::prepare(problem, std::move(materialComponent));
            preconditioning::prepare(
                problem.GetPhysicalOperator().GetHydrostaticOperator().GetFEM(),
                problem.GetPhysicalOperator().GetGravityContext().GetGeometryContext(), std::move(gravityComponent)
            );
            StellarStructureCrossJacobianOperator{problem.GetPhysicalOperator()};
        };

    template <
        equilibrium::StellarEquilibriumModel Model,
        equilibrium::StellarDiscretizationType Discretization,
        typename MaterialComponent,
        backend::Registered GravityMassBackend,
        backend::ApplicationMode Mode,
        GravityFactorizationPolicy GravityPolicy,
        StellarStructureFactorizationPolicy StructurePolicy>
        requires StellarStructurePreparableFor<
            equilibrium::StellarEquilibriumProblem<Model, Discretization>,
            MaterialComponent,
            GravityFieldBlock<GravityMassBackend, backend::HypreBoomerAMG<Mode>, GravityPolicy>>
    class PreparedStellarStructureBlock final : public mfem::Solver {
    private:
        using Problem          = equilibrium::StellarEquilibriumProblem<Model, Discretization>;
        using GravityComponent = GravityFieldBlock<GravityMassBackend, backend::HypreBoomerAMG<Mode>, GravityPolicy>;
        using Structure        = StellarStructureBlock<
            MaterialComponent,
            GravityComponent,
            typename Problem::FormType,
            typename Problem::JacobianFormType,
            StructurePolicy>;
        using MaterialPrepared =
            decltype(preconditioning::prepare(std::declval<const Problem &>(), std::declval<MaterialComponent>()));
        using GravityPrepared = decltype(preconditioning::prepare(
            std::declval<const fem::FEM &>(),
            std::declval<const operators::context::gravity_field::GravityFieldGeometryContext &>(),
            std::declval<GravityComponent>()
        ));

    public:
        PreparedStellarStructureBlock(
            const Problem &problem,
            Structure structure
        )
            : mfem::Solver(StructureSize(problem)),
              m_problem(std::addressof(problem)),
              m_structure(std::move(structure)),
              m_materialSurface(
                  preconditioning::prepare(
                      problem,
                      m_structure.materialSurfaceComponent()
                  )
              ),
              m_gravity(
                  preconditioning::prepare(
                      detail::physicalOperator(problem).GetHydrostaticOperator().GetFEM(),
                      detail::physicalOperator(problem).GetGravityContext().GetGeometryContext(),
                      m_structure.gravityComponent()
                  )
              ),
              m_crossCouplings(detail::physicalOperator(problem)),
              m_factorization(
                  m_structure.factorizationPolicy(),
                  m_materialSurface,
                  m_gravity,
                  m_crossCouplings
              ) {
        }

        PreparedStellarStructureBlock(const PreparedStellarStructureBlock &)            = delete;
        PreparedStellarStructureBlock &operator=(const PreparedStellarStructureBlock &) = delete;
        PreparedStellarStructureBlock(PreparedStellarStructureBlock &&)                 = delete;
        PreparedStellarStructureBlock &operator=(PreparedStellarStructureBlock &&)      = 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 stellar-structure block is stale; refresh it before application.");
            }
            m_factorization.Mult(rightHandSide, action);
        }

        [[nodiscard]] bool IsCurrent() const noexcept {
            return m_materialSurface.IsCurrent() && m_gravity.IsCurrent() &&
                   detail::physicalOperator(*m_problem).IsPrepared();
        }

        [[nodiscard]] StellarStructureBlockPreparationReport Refresh() {
            const auto &physical = detail::physicalOperator(*m_problem);
            return {
                .materialSurface = m_materialSurface.Refresh(physical),
                .gravity         = m_gravity.Refresh(
                    physical.GetHydrostaticOperator().GetFEM(), physical.GetGravityContext().GetGeometryContext()
                )
            };
        }

        [[nodiscard]] const Structure &GetBlock() const noexcept {
            return m_structure;
        }

        [[nodiscard]] const MaterialPrepared &GetMaterialSurfacePreconditioner() const noexcept {
            return m_materialSurface;
        }

        [[nodiscard]] const GravityPrepared &GetGravityPreconditioner() const noexcept {
            return m_gravity;
        }

        [[nodiscard]] const StellarStructureCrossJacobianOperator &GetCrossCouplings() const noexcept {
            return m_crossCouplings;
        }

        [[nodiscard]] const StellarStructureFactorizationOperator<
            StructurePolicy,
            StellarStructureCrossJacobianOperator> &
        GetFactorization() const noexcept {
            return m_factorization;
        }

    private:
        [[nodiscard]] static int StructureSize(const Problem &problem) {
            const auto &physical = detail::physicalOperator(problem);
            return physical.GetGravityContext().GetDensityMap().reduced_size() +
                   physical.GetDomainDeformation().parameterCount() +
                   physical.GetBarotropicClosureOperator().GetEnthalpySize() +
                   physical.GetGravityContext().GetGravityGradientMap().reduced_size() +
                   physical.GetGravityContext().GetGravityPotentialMap().reduced_size();
        }

        const Problem *m_problem;
        Structure m_structure;
        MaterialPrepared m_materialSurface;
        GravityPrepared m_gravity;
        StellarStructureCrossJacobianOperator m_crossCouplings;
        StellarStructureFactorizationOperator<StructurePolicy, StellarStructureCrossJacobianOperator> m_factorization;
    };

    template <
        StellarStructurePreconditionerProblem Problem,
        typename MaterialComponent,
        backend::Registered GravityMassBackend,
        backend::ApplicationMode Mode,
        GravityFactorizationPolicy GravityPolicy,
        StellarStructureFactorizationPolicy StructurePolicy>
    [[nodiscard]] constexpr auto stellarStructureBlock(
        const Problem &,
        MaterialComponent materialComponent,
        GravityFieldBlock<
            GravityMassBackend,
            backend::HypreBoomerAMG<Mode>,
            GravityPolicy> gravityComponent,
        StructurePolicy policy
    ) {
        using ProblemType = std::remove_cvref_t<Problem>;
        return StellarStructureBlock<
            MaterialComponent, GravityFieldBlock<GravityMassBackend, backend::HypreBoomerAMG<Mode>, GravityPolicy>,
            typename ProblemType::FormType, typename ProblemType::JacobianFormType, StructurePolicy>{
            std::move(materialComponent), std::move(gravityComponent), std::move(policy)
        };
    }

    template <StellarStructurePreconditionerProblem Problem>
    [[nodiscard]] constexpr auto stellarStructureBlock(const Problem &problem) {
        using FixedAMG = backend::HypreBoomerAMG<backend::FixedCycles>;
        auto material  = materialSurfaceBlock(problem);
        auto gravity   = GravityFieldBlock(
            backend::MatrixFreeChebyshev{.order = 5, .powerIterations = 20},
            FixedAMG{backend::FixedCycles{.cycles = 3}}, GravityApproximateLDU{}
        );
        return stellarStructureBlock(problem, std::move(material), std::move(gravity), IndependentStellarSubsystems{});
    }

    template <
        equilibrium::StellarEquilibriumModel Model,
        equilibrium::StellarDiscretizationType Discretization,
        typename MaterialComponent,
        backend::Registered GravityMassBackend,
        backend::ApplicationMode Mode,
        GravityFactorizationPolicy GravityPolicy,
        StellarStructureFactorizationPolicy StructurePolicy>
        requires StellarStructurePreparableFor<
            equilibrium::StellarEquilibriumProblem<
                Model,
                Discretization>,
            MaterialComponent,
            GravityFieldBlock<
                GravityMassBackend,
                backend::HypreBoomerAMG<Mode>,
                GravityPolicy>>
    [[nodiscard]] auto prepare(
        const equilibrium::StellarEquilibriumProblem<
            Model,
            Discretization> &problem,
        StellarStructureBlock<
            MaterialComponent,
            GravityFieldBlock<
                GravityMassBackend,
                backend::HypreBoomerAMG<Mode>,
                GravityPolicy>,
            typename equilibrium::StellarEquilibriumProblem<
                Model,
                Discretization>::FormType,
            typename equilibrium::StellarEquilibriumProblem<
                Model,
                Discretization>::JacobianFormType,
            StructurePolicy> structure
    ) {
        return PreparedStellarStructureBlock<
            Model, Discretization, MaterialComponent, GravityMassBackend, Mode, GravityPolicy, StructurePolicy>{
            problem, std::move(structure)
        };
    }
} // namespace mean_field::preconditioning