module;

#include <algorithm>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <stdexcept>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>

#include <mfem.hpp>

export module mean_field:preconditioning.specification_border;

export import :preconditioning.stellar_equilibrium;
export import :preconditioning.stellar_structure;

export namespace mean_field::preconditioning {
    template <models::ModelSpecification Specification> struct SpecificationBorderContribution {
        using CorrectionBlocks           = utils::blocks::type_list<>;
        using ResidualBlocks             = utils::blocks::type_list<>;
        using RequiredCouplings          = utils::blocks::type_list<>;

        static constexpr bool registered = false;
    };

    template <> struct SpecificationBorderContribution<models::FixedTotalMass> {
        using LayoutRequest     = models::FixedMassLayoutRequest;
        using CorrectionBlock   = typename LayoutRequest::ValueBlockType;
        using ResidualBlock     = typename LayoutRequest::ResidualBlockType;
        using CorrectionBlocks  = utils::blocks::type_list<CorrectionBlock>;
        using ResidualBlocks    = utils::blocks::type_list<ResidualBlock>;
        using RequiredCouplings = utils::blocks::type_list<
            Coupling<utils::blocks::enthalpy::specific::residual, CorrectionBlock>,
            Coupling<ResidualBlock, utils::blocks::density::mass::value>,
            Coupling<ResidualBlock, utils::blocks::surface_deformation::parameters::value>>;

        static constexpr bool registered = true;
    };

    template <> struct SpecificationBorderContribution<models::FixedCentralDensity> {
        using LayoutRequest     = models::CentralDensityLayoutRequest;
        using CorrectionBlock   = typename LayoutRequest::ValueBlockType;
        using ResidualBlock     = typename LayoutRequest::ResidualBlockType;
        using CorrectionBlocks  = utils::blocks::type_list<CorrectionBlock>;
        using ResidualBlocks    = utils::blocks::type_list<ResidualBlock>;
        using RequiredCouplings = utils::blocks::type_list<
            Coupling<utils::blocks::enthalpy::specific::residual, CorrectionBlock>,
            Coupling<ResidualBlock, utils::blocks::enthalpy::specific::value>>;

        static constexpr bool registered = true;
    };

    namespace detail {
        template <models::ModelSpecification Specification>
        inline constexpr std::size_t generatedBorderValueArity =
            models::specificationDescriptor<Specification>().generatedValueArity;

        template <models::ModelSpecification Specification>
        inline constexpr std::size_t generatedBorderResidualArity =
            models::specificationDescriptor<Specification>().generatedResidualArity;

        template <models::ModelSpecification Specification>
        inline constexpr bool specificationGeneratesBorder =
            generatedBorderValueArity<Specification> != 0 || generatedBorderResidualArity<Specification> != 0;

        template <models::ModelSpecification Specification>
        inline constexpr bool specificationBorderContributionIsComplete =
            !specificationGeneratesBorder<Specification> ||
            (SpecificationBorderContribution<Specification>::registered &&
             generatedBorderValueArity<Specification> == generatedBorderResidualArity<Specification> &&
             SpecificationBorderContribution<Specification>::CorrectionBlocks::size == 1 &&
             SpecificationBorderContribution<Specification>::ResidualBlocks::size == 1);

        template <typename SpecificationSet> struct CompiledSpecificationBorder;

        template <models::ModelSpecification... Specifications>
        struct CompiledSpecificationBorder<models::detail::SpecificationSetStorage<Specifications...>> {
            static_assert(
                (specificationBorderContributionIsComplete<Specifications> && ...),
                "Every specification-generated border requires a registered preconditioning contribution with "
                "balanced value and residual arity."
            );

            using SpecificationTypes = models::detail::SpecificationSetStorage<Specifications...>;
            using CorrectionBlocks   = preconditioning::detail::ConcatenateT<
                typename SpecificationBorderContribution<Specifications>::CorrectionBlocks...>;
            using ResidualBlocks = preconditioning::detail::ConcatenateT<
                typename SpecificationBorderContribution<Specifications>::ResidualBlocks...>;
            using RequiredCouplings = preconditioning::detail::ConcatenateT<
                typename SpecificationBorderContribution<Specifications>::RequiredCouplings...>;

            static constexpr std::size_t valueArity =
                (std::size_t{0} + ... + generatedBorderValueArity<Specifications>);
            static constexpr std::size_t residualArity =
                (std::size_t{0} + ... + generatedBorderResidualArity<Specifications>);
            static constexpr std::size_t specificationCount =
                (std::size_t{0} + ... + (SpecificationBorderContribution<Specifications>::registered ? 1U : 0U));
            static constexpr bool symbolicallySquare = valueArity == residualArity;
        };

        template <typename Query, typename SpecificationSet> struct SpecificationBorderValueOffset;

        template <typename Query, models::ModelSpecification Head, models::ModelSpecification... Tail>
        struct SpecificationBorderValueOffset<Query, models::detail::SpecificationSetStorage<Head, Tail...>> {
            static constexpr std::size_t value = [] {
                if constexpr (std::same_as<Query, Head>) {
                    return std::size_t{0};
                } else {
                    static_assert(sizeof...(Tail) > 0, "The requested border specification is not in the model.");
                    return generatedBorderValueArity<Head> +
                           SpecificationBorderValueOffset<
                               Query, models::detail::SpecificationSetStorage<Tail...>>::value;
                }
            }();
        };

        template <typename Query, typename SpecificationSet> struct SpecificationBorderResidualOffset;

        template <typename Query, models::ModelSpecification Head, models::ModelSpecification... Tail>
        struct SpecificationBorderResidualOffset<Query, models::detail::SpecificationSetStorage<Head, Tail...>> {
            static constexpr std::size_t value = [] {
                if constexpr (std::same_as<Query, Head>) {
                    return std::size_t{0};
                } else {
                    static_assert(sizeof...(Tail) > 0, "The requested border specification is not in the model.");
                    return generatedBorderResidualArity<Head> +
                           SpecificationBorderResidualOffset<
                               Query, models::detail::SpecificationSetStorage<Tail...>>::value;
                }
            }();
        };
    } // namespace detail

    template <model::StellarModelType Model>
    using CompiledSpecificationBorderFor =
        detail::CompiledSpecificationBorder<typename std::remove_cvref_t<Model>::SpecificationTypes>;

    template <models::ModelSpecification Specification, model::StellarModelType Model>
    inline constexpr std::size_t specificationBorderValueOffset = detail::
        SpecificationBorderValueOffset<Specification, typename std::remove_cvref_t<Model>::SpecificationTypes>::value;

    template <models::ModelSpecification Specification, model::StellarModelType Model>
    inline constexpr std::size_t specificationBorderResidualOffset = detail::SpecificationBorderResidualOffset<
        Specification,
        typename std::remove_cvref_t<Model>::SpecificationTypes>::value;

    using SpecificationBorderCharacteristics = OperatorCharacteristics<
        OperatorCategory::dense_border,
        OperatorValueStructure::block,
        OperatorSymmetry::nonsymmetric,
        OperatorDefiniteness::indefinite,
        OperatorRepresentation::assembled_dense,
        OperatorDistribution::local>;

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

    namespace backend {
        template <Registered StructureBackend, Registered BorderBackend = DenseDirect>
        struct BorderedStellarStructure final {
            using StructureBackendType = StructureBackend;
            using BorderBackendType    = BorderBackend;
        };

        template <Registered StructureBackend, Registered BorderBackend>
        struct Traits<BorderedStellarStructure<StructureBackend, BorderBackend>> {
            static constexpr bool registered = true;
            static constexpr ApplicationContract applicationContract =
                ::mean_field::preconditioning::backend::applicationContract<StructureBackend> ==
                            ApplicationContract::stationary_linear &&
                        ::mean_field::preconditioning::backend::applicationContract<BorderBackend> ==
                            ApplicationContract::stationary_linear
                    ? ApplicationContract::stationary_linear
                    : ApplicationContract::flexible;
            static constexpr bool supportsSerialExecution = Traits<StructureBackend>::supportsSerialExecution;
            static constexpr bool supportsDistributedExecution =
                Traits<StructureBackend>::supportsDistributedExecution &&
                Traits<BorderBackend>::supportsSerialExecution;
            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<StructureBackend>::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 StructureComponentT,
        model::StellarModelType ModelT,
        typename FormT,
        typename JacobianFormT>
        requires utils::blocks::valid_jacobian_form<FormT, JacobianFormT>
    class SpecificationBorderBlock final {
    private:
        using CompiledBorder = CompiledSpecificationBorderFor<ModelT>;

    public:
        using StructureComponent = StructureComponentT;
        using Model              = ModelT;
        using Form               = FormT;
        using JacobianForm       = JacobianFormT;
        using CorrectionBlocks   = preconditioning::detail::
            ConcatenateT<typename StructureComponent::CorrectionBlocks, typename CompiledBorder::CorrectionBlocks>;
        using ResidualBlocks = preconditioning::detail::
            ConcatenateT<typename StructureComponent::ResidualBlocks, typename CompiledBorder::ResidualBlocks>;
        using RequiredCouplings = preconditioning::detail::
            ConcatenateT<typename StructureComponent::RequiredCouplings, typename CompiledBorder::RequiredCouplings>;
        using OperatorDescription = BorderedStellarStructureCharacteristics;
        using BackendType =
            backend::BorderedStellarStructure<typename StructureComponent::BackendType, backend::DenseDirect>;
        using PreparationDependencies                 = typename backend::Traits<BackendType>::PreparationDependencies;

        static constexpr std::size_t borderValueArity = CompiledBorder::valueArity;
        static constexpr std::size_t borderResidualArity = CompiledBorder::residualArity;

        constexpr explicit SpecificationBorderBlock(
            StructureComponent structureComponent,
            backend::DenseDirect borderBackend = {}
        )
            : m_structureComponent(std::move(structureComponent)),
              m_borderBackend(std::move(borderBackend)) {
            static_assert(CompiledBorder::symbolicallySquare);
        }

        [[nodiscard]] constexpr const StructureComponent &structureComponent() const noexcept {
            return m_structureComponent;
        }

        [[nodiscard]] constexpr const backend::DenseDirect &borderBackend() const noexcept {
            return m_borderBackend;
        }

    private:
        StructureComponent m_structureComponent;
        backend::DenseDirect m_borderBackend;
    };

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

    template <
        PreconditionerComponent StructureComponent,
        model::StellarModelType Model,
        typename Form,
        typename JacobianForm>
    struct IsSpecificationBorderBlock<SpecificationBorderBlock<StructureComponent, Model, Form, JacobianForm>>
        : std::true_type { };

    template <typename Candidate>
    concept SpecificationBorderBlockType = IsSpecificationBorderBlock<std::remove_cvref_t<Candidate>>::value;

    struct StellarStructureDirectionView final {
        const mfem::Vector &density;
        const mfem::Vector &surface;
        const mfem::Vector &enthalpy;
        const mfem::Vector &gravityGradient;
        const mfem::Vector &gravityPotential;
    };

    struct StellarStructureActionView final {
        mfem::Vector &density;
        mfem::Vector &surface;
        mfem::Vector &enthalpy;
        mfem::Vector &gravityGradient;
        mfem::Vector &gravityPotential;
    };

    namespace detail {
        template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        [[nodiscard]] const operators::PreparedStellarEquilibriumOperator &
        specificationBorderPhysicalOperator(const Problem &problem) {
            if constexpr (std::remove_cvref_t<Problem>::hasFixedCentralDensity) {
                return problem.GetPreparedOperator().GetPhysicalOperator();
            } else {
                return problem.GetPreparedOperator();
            }
        }

        template <models::ModelSpecification Specification, equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        class PreparedSpecificationBorderAction {
            static_assert(
                !specificationGeneratesBorder<Specification>,
                "A generated model specification requires a prepared specification-border action specialization."
            );

        public:
            explicit PreparedSpecificationBorderAction(const Problem &) noexcept {
            }

            void ApplyStructureToBorder(
                const StellarStructureDirectionView &,
                mfem::Vector &
            ) const noexcept {
            }

            void ApplyBorderToStructure(
                const mfem::Vector &,
                StellarStructureActionView
            ) const noexcept {
            }

            void ApplyBorderToBorder(
                const mfem::Vector &,
                mfem::Vector &
            ) const noexcept {
            }
        };

        template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        class PreparedSpecificationBorderAction<models::FixedTotalMass, Problem> {
        private:
            using Model = typename std::remove_cvref_t<Problem>::ModelType;

        public:
            explicit PreparedSpecificationBorderAction(const Problem &problem)
                : m_physical(std::addressof(specificationBorderPhysicalOperator(problem))),
                  m_volumeDisplacement(m_physical->GetDomainDeformation().volumeDisplacementSize()),
                  m_enthalpyWorkspace(m_physical->GetBarotropicClosureOperator().GetEnthalpySize()),
                  m_zeroEnthalpy(m_physical->GetBarotropicClosureOperator().GetEnthalpySize()) {
                m_zeroEnthalpy = 0.0;
            }

            void ApplyStructureToBorder(
                const StellarStructureDirectionView &structure,
                mfem::Vector &borderAction
            ) const {
                constexpr int residualOffset =
                    static_cast<int>(specificationBorderResidualOffset<models::FixedTotalMass, Model>);
                mfem::Vector massAction(borderAction, residualOffset, 1);
                m_physical->GetDomainDeformation().applyJacobian(
                    m_physical->GetSurfaceDeformationParameters(), structure.surface, m_volumeDisplacement
                );
                m_physical->GetMassNormalizationOperator().ApplyCompleteJacobianAction(
                    structure.density, m_volumeDisplacement, massAction
                );
                massAction.SyncAliasMemory(borderAction);
            }

            void ApplyBorderToStructure(
                const mfem::Vector &borderDirection,
                StellarStructureActionView structureAction
            ) const {
                constexpr int valueOffset =
                    static_cast<int>(specificationBorderValueOffset<models::FixedTotalMass, Model>);
                m_physical->GetHydrostaticOperator().ApplyBernoulliConstantJacobianAction(
                    borderDirection(valueOffset), m_enthalpyWorkspace
                );
                m_physical->GetSurfaceConstraintOperator().ApplyJacobianRows(m_zeroEnthalpy, m_enthalpyWorkspace);
                structureAction.enthalpy += m_enthalpyWorkspace;
            }

            void ApplyBorderToBorder(
                const mfem::Vector &,
                mfem::Vector &
            ) const noexcept {
            }

        private:
            const operators::PreparedStellarEquilibriumOperator *m_physical;
            mutable mfem::Vector m_volumeDisplacement;
            mutable mfem::Vector m_enthalpyWorkspace;
            mfem::Vector m_zeroEnthalpy;
        };

        template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        class PreparedSpecificationBorderAction<models::FixedCentralDensity, Problem> {
        private:
            using ProblemType = std::remove_cvref_t<Problem>;
            using Model       = typename ProblemType::ModelType;
            static_assert(ProblemType::hasFixedCentralDensity);

        public:
            explicit PreparedSpecificationBorderAction(const Problem &problem)
                : m_constraint(std::addressof(problem.GetPreparedOperator().GetCentralDensityConstraint())),
                  m_zeroEnthalpy(
                      specificationBorderPhysicalOperator(problem).GetBarotropicClosureOperator().GetEnthalpySize()
                  ),
                  m_enthalpyWorkspace(m_zeroEnthalpy.Size()),
                  m_phaseWorkspace(1) {
                m_zeroEnthalpy = 0.0;
            }

            void ApplyStructureToBorder(
                const StellarStructureDirectionView &structure,
                mfem::Vector &borderAction
            ) const {
                constexpr int residualOffset =
                    static_cast<int>(specificationBorderResidualOffset<models::FixedCentralDensity, Model>);
                mfem::Vector phaseAction(borderAction, residualOffset, 1);
                m_enthalpyWorkspace = 0.0;
                m_constraint->ApplyJacobian(
                    {.enthalpyVariation = structure.enthalpy, .borderVariation = 0.0},
                    {.enthalpyAction = m_enthalpyWorkspace, .phaseAction = phaseAction}
                );
                phaseAction.SyncAliasMemory(borderAction);
            }

            void ApplyBorderToStructure(
                const mfem::Vector &borderDirection,
                StellarStructureActionView structureAction
            ) const {
                constexpr int valueOffset =
                    static_cast<int>(specificationBorderValueOffset<models::FixedCentralDensity, Model>);
                m_enthalpyWorkspace = 0.0;
                m_phaseWorkspace    = 0.0;
                m_constraint->ApplyJacobian(
                    {.enthalpyVariation = m_zeroEnthalpy, .borderVariation = borderDirection(valueOffset)},
                    {.enthalpyAction = m_enthalpyWorkspace, .phaseAction = m_phaseWorkspace}
                );
                structureAction.enthalpy += m_enthalpyWorkspace;
            }

            void ApplyBorderToBorder(
                const mfem::Vector &,
                mfem::Vector &
            ) const noexcept {
            }

        private:
            const operators::PreparedCentralDensityConstraint *m_constraint;
            mfem::Vector m_zeroEnthalpy;
            mutable mfem::Vector m_enthalpyWorkspace;
            mutable mfem::Vector m_phaseWorkspace;
        };

        template <typename SpecificationSet, equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        class PreparedSpecificationBorderActions;

        template <
            models::ModelSpecification... Specifications,
            equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        class PreparedSpecificationBorderActions<models::detail::SpecificationSetStorage<Specifications...>, Problem> {
        public:
            explicit PreparedSpecificationBorderActions(const Problem &problem)
                : m_actions(
                      PreparedSpecificationBorderAction<
                          Specifications,
                          Problem>{problem}...
                  ) {
            }

            void ApplyStructureToBorder(
                const StellarStructureDirectionView &structure,
                mfem::Vector &borderAction
            ) const {
                std::apply(
                    [&](const auto &...actions) { (actions.ApplyStructureToBorder(structure, borderAction), ...); },
                    m_actions
                );
            }

            void ApplyBorderToStructure(
                const mfem::Vector &borderDirection,
                StellarStructureActionView structureAction
            ) const {
                std::apply(
                    [&](const auto &...actions) {
                        (actions.ApplyBorderToStructure(borderDirection, structureAction), ...);
                    },
                    m_actions
                );
            }

            void ApplyBorderToBorder(
                const mfem::Vector &borderDirection,
                mfem::Vector &borderAction
            ) const {
                std::apply(
                    [&](const auto &...actions) { (actions.ApplyBorderToBorder(borderDirection, borderAction), ...); },
                    m_actions
                );
            }

        private:
            std::tuple<PreparedSpecificationBorderAction<Specifications, Problem>...> m_actions;
        };
    } // namespace detail

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
    class SpecificationBorderJacobianOperator final : public mfem::Operator {
    private:
        using ProblemType    = std::remove_cvref_t<Problem>;
        using Model          = typename ProblemType::ModelType;
        using CompiledBorder = CompiledSpecificationBorderFor<Model>;
        using Actions = detail::PreparedSpecificationBorderActions<typename Model::SpecificationTypes, ProblemType>;

    public:
        explicit SpecificationBorderJacobianOperator(const ProblemType &problem)
            : mfem::Operator(StructureSizeOf(problem) + BorderSizeOf(problem)),
              m_structureOffsets(6),
              m_actions(problem) {
            const auto &physical  = detail::specificationBorderPhysicalOperator(problem);
            m_structureOffsets[0] = 0;
            m_structureOffsets[1] = physical.GetGravityContext().GetDensityMap().reduced_size();
            m_structureOffsets[2] = m_structureOffsets[1] + physical.GetDomainDeformation().parameterCount();
            m_structureOffsets[3] = m_structureOffsets[2] + physical.GetBarotropicClosureOperator().GetEnthalpySize();
            m_structureOffsets[4] =
                m_structureOffsets[3] + physical.GetGravityContext().GetGravityGradientMap().reduced_size();
            m_structureOffsets[5] = StructureSizeOf(problem);

            if (StructureSize() + BorderSize() != problem.StateSize() ||
                StructureSize() + BorderSize() != problem.EquationSize()) {
                throw std::logic_error(
                    "The compiled specification border does not complete the stellar-equilibrium problem."
                );
            }
        }

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override {
            VerifyCombined(direction, action);
            action = 0.0;
            const mfem::Vector structureDirection(const_cast<mfem::real_t *>(direction.GetData()), StructureSize());
            const mfem::Vector borderDirection(
                const_cast<mfem::real_t *>(direction.GetData()) + StructureSize(), BorderSize()
            );
            mfem::Vector structureAction(action, 0, StructureSize());
            mfem::Vector borderAction(action, StructureSize(), BorderSize());

            ApplyBorderToStructure(borderDirection, structureAction);
            ApplyStructureToBorder(structureDirection, borderAction);
            mfem::Vector borderDiagonalAction(BorderSize());
            ApplyBorderToBorder(borderDirection, borderDiagonalAction);
            borderAction += borderDiagonalAction;

            structureAction.SyncAliasMemory(action);
            borderAction.SyncAliasMemory(action);
        }

        void ApplyStructureToBorder(
            const mfem::Vector &structureDirection,
            mfem::Vector &borderAction
        ) const {
            VerifyStructure(structureDirection, "direction");
            VerifyBorder(borderAction, "action");
            borderAction             = 0.0;
            const auto directionView = StructureDirection(structureDirection);
            m_actions.ApplyStructureToBorder(directionView, borderAction);
        }

        void ApplyBorderToStructure(
            const mfem::Vector &borderDirection,
            mfem::Vector &structureAction
        ) const {
            VerifyBorder(borderDirection, "direction");
            VerifyStructure(structureAction, "action");
            structureAction             = 0.0;
            auto densityAction          = MutableStructureBlock(structureAction, 0);
            auto surfaceAction          = MutableStructureBlock(structureAction, 1);
            auto enthalpyAction         = MutableStructureBlock(structureAction, 2);
            auto gravityGradientAction  = MutableStructureBlock(structureAction, 3);
            auto gravityPotentialAction = MutableStructureBlock(structureAction, 4);
            m_actions.ApplyBorderToStructure(
                borderDirection, {.density          = densityAction,
                                  .surface          = surfaceAction,
                                  .enthalpy         = enthalpyAction,
                                  .gravityGradient  = gravityGradientAction,
                                  .gravityPotential = gravityPotentialAction}
            );
            densityAction.SyncAliasMemory(structureAction);
            surfaceAction.SyncAliasMemory(structureAction);
            enthalpyAction.SyncAliasMemory(structureAction);
            gravityGradientAction.SyncAliasMemory(structureAction);
            gravityPotentialAction.SyncAliasMemory(structureAction);
        }

        void ApplyBorderToBorder(
            const mfem::Vector &borderDirection,
            mfem::Vector &borderAction
        ) const {
            VerifyBorder(borderDirection, "direction");
            VerifyBorder(borderAction, "action");
            borderAction = 0.0;
            m_actions.ApplyBorderToBorder(borderDirection, borderAction);
        }

        [[nodiscard]] int StructureSize() const noexcept {
            return m_structureOffsets.Last();
        }

        [[nodiscard]] static constexpr int BorderSize() noexcept {
            return static_cast<int>(CompiledBorder::valueArity);
        }

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

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

        [[nodiscard]] static constexpr int BorderSizeOf(const ProblemType &) noexcept {
            return BorderSize();
        }

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

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

        [[nodiscard]] StellarStructureDirectionView StructureDirection(const mfem::Vector &direction) const {
            m_directionDensity          = ConstStructureBlock(direction, 0);
            m_directionSurface          = ConstStructureBlock(direction, 1);
            m_directionEnthalpy         = ConstStructureBlock(direction, 2);
            m_directionGravityGradient  = ConstStructureBlock(direction, 3);
            m_directionGravityPotential = ConstStructureBlock(direction, 4);
            return {
                .density          = m_directionDensity,
                .surface          = m_directionSurface,
                .enthalpy         = m_directionEnthalpy,
                .gravityGradient  = m_directionGravityGradient,
                .gravityPotential = m_directionGravityPotential
            };
        }

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

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

        void VerifyBorder(
            const mfem::Vector &vector,
            const char *role
        ) const {
            if (vector.Size() != BorderSize()) {
                throw std::invalid_argument(std::string("The specification border ") + role + " has the wrong size.");
            }
        }

        mfem::Array<int> m_structureOffsets;
        Actions m_actions;
        mutable mfem::Vector m_directionDensity;
        mutable mfem::Vector m_directionSurface;
        mutable mfem::Vector m_directionEnthalpy;
        mutable mfem::Vector m_directionGravityGradient;
        mutable mfem::Vector m_directionGravityPotential;
    };

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
    SpecificationBorderJacobianOperator(const Problem &)
        -> SpecificationBorderJacobianOperator<std::remove_cvref_t<Problem>>;

    template <typename Candidate>
    concept SpecificationBorderCouplingOperator = requires(
        const Candidate &couplings,
        const mfem::Vector &structureDirection,
        const mfem::Vector &borderDirection,
        mfem::Vector &structureAction,
        mfem::Vector &borderAction
    ) {
        { couplings.StructureSize() } -> std::same_as<int>;
        { couplings.BorderSize() } -> std::same_as<int>;
        couplings.ApplyStructureToBorder(structureDirection, borderAction);
        couplings.ApplyBorderToStructure(borderDirection, structureAction);
        couplings.ApplyBorderToBorder(borderDirection, borderAction);
    };

    struct SpecificationBorderFactorizationStatistics final {
        std::uint64_t setups{0};
        std::uint64_t applications{0};
        std::uint64_t structureInverseApplications{0};
        std::uint64_t cachedStructureInverseBorderApplications{0};
        std::uint64_t structureToBorderApplications{0};
        std::uint64_t borderToStructureApplications{0};
        std::uint64_t borderToBorderApplications{0};
        std::uint64_t schurProbes{0};
    };

    template <
        SpecificationBorderCouplingOperator CouplingOperator,
        ApplicationContract StructureInverseContract = ApplicationContract::stationary_linear>
    class SpecificationBorderFactorizationOperator final : public mfem::Solver {
    public:
        static constexpr bool cachesStructureInverseBorderCoupling =
            StructureInverseContract == ApplicationContract::stationary_linear;

        SpecificationBorderFactorizationOperator(
            const mfem::Solver &structureInverse,
            const CouplingOperator &couplings,
            backend::DenseDirect borderBackend = {}
        )
            : mfem::Solver(couplings.StructureSize() + couplings.BorderSize()),
              m_structureInverse(std::addressof(structureInverse)),
              m_couplings(std::addressof(couplings)),
              m_borderBackend(std::move(borderBackend)),
              m_schurComplement(couplings.BorderSize()),
              m_structureInverseBorderCoupling(
                  cachesStructureInverseBorderCoupling ? couplings.StructureSize() : 0,
                  cachesStructureInverseBorderCoupling ? couplings.BorderSize() : 0
              ),
              m_structureWorkspace(couplings.StructureSize()),
              m_structureCoupling(couplings.StructureSize()),
              m_borderWorkspace(couplings.BorderSize()),
              m_borderCoupling(couplings.BorderSize()),
              m_borderDiagonal(couplings.BorderSize()),
              m_borderBasis(couplings.BorderSize()) {
            if (structureInverse.Height() <= 0 || structureInverse.Height() != structureInverse.Width() ||
                structureInverse.Height() != couplings.StructureSize() || couplings.BorderSize() < 0) {
                throw std::invalid_argument(
                    "The structure inverse and specification-border couplings have incompatible dimensions."
                );
            }
            AssembleSchurComplement();
        }

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

        void SetOperator(const mfem::Operator &operation) override {
            if (operation.Height() != Height() || operation.Width() != Width()) {
                throw std::invalid_argument(
                    "The specification-border 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 specification-border factorization requires compatible, preallocated vectors."
                );
            }

            const mfem::Vector structureRightHandSide(
                const_cast<mfem::real_t *>(rightHandSide.GetData()), m_couplings->StructureSize()
            );
            const mfem::Vector borderRightHandSide(
                const_cast<mfem::real_t *>(rightHandSide.GetData()) + m_couplings->StructureSize(),
                m_couplings->BorderSize()
            );
            action = 0.0;
            mfem::Vector structureAction(action, 0, m_couplings->StructureSize());
            mfem::Vector borderAction(action, m_couplings->StructureSize(), m_couplings->BorderSize());

            if (m_couplings->BorderSize() == 0) {
                m_structureInverse->Mult(structureRightHandSide, structureAction);
                ++m_statistics.structureInverseApplications;
            } else {
                m_structureInverse->Mult(structureRightHandSide, m_structureWorkspace);
                m_couplings->ApplyStructureToBorder(m_structureWorkspace, m_borderCoupling);
                m_borderWorkspace = borderRightHandSide;
                m_borderWorkspace -= m_borderCoupling;
                m_borderInverse->Mult(m_borderWorkspace, borderAction);

                if constexpr (cachesStructureInverseBorderCoupling) {
                    // W = A^{-1} B was assembled with the border Schur complement, so the
                    // stationary-linear structure correction is A^{-1} f - W y.
                    m_structureInverseBorderCoupling.Mult(borderAction, m_structureCoupling);
                    structureAction = m_structureWorkspace;
                    structureAction -= m_structureCoupling;
                    ++m_statistics.structureInverseApplications;
                    ++m_statistics.cachedStructureInverseBorderApplications;
                } else {
                    m_couplings->ApplyBorderToStructure(borderAction, m_structureCoupling);
                    m_structureCoupling *= -1.0;
                    m_structureCoupling += structureRightHandSide;
                    m_structureInverse->Mult(m_structureCoupling, structureAction);
                    m_statistics.structureInverseApplications += 2;
                    ++m_statistics.borderToStructureApplications;
                }
                ++m_statistics.structureToBorderApplications;
            }

            structureAction.SyncAliasMemory(action);
            borderAction.SyncAliasMemory(action);
            ++m_statistics.applications;
        }

        void RefreshSchurComplement() {
            AssembleSchurComplement();
        }

        [[nodiscard]] const mfem::DenseMatrix &GetSchurComplement() const noexcept {
            return m_schurComplement;
        }

        [[nodiscard]] const backend::PreparedDenseDirect *GetBorderInverse() const noexcept {
            return m_borderInverse.get();
        }

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

    private:
        void AssembleSchurComplement() {
            const int borderSize = m_couplings->BorderSize();
            if (borderSize == 0) {
                m_schurComplement.SetSize(0, 0);
                m_borderInverse.reset();
                ++m_statistics.setups;
                return;
            }

            mfem::Vector schurColumn(borderSize);
            for (int column = 0; column < borderSize; ++column) {
                m_borderBasis         = 0.0;
                m_borderBasis(column) = 1.0;
                m_couplings->ApplyBorderToStructure(m_borderBasis, m_structureCoupling);
                ++m_statistics.borderToStructureApplications;
                m_structureInverse->Mult(m_structureCoupling, m_structureWorkspace);
                ++m_statistics.structureInverseApplications;
                if constexpr (cachesStructureInverseBorderCoupling) {
                    m_structureInverseBorderCoupling.SetCol(column, m_structureWorkspace);
                }
                m_couplings->ApplyStructureToBorder(m_structureWorkspace, m_borderCoupling);
                ++m_statistics.structureToBorderApplications;
                m_couplings->ApplyBorderToBorder(m_borderBasis, m_borderDiagonal);
                ++m_statistics.borderToBorderApplications;
                schurColumn = m_borderDiagonal;
                schurColumn -= m_borderCoupling;
                for (int row = 0; row < borderSize; ++row) {
                    m_schurComplement(row, column) = schurColumn(row);
                }
                ++m_statistics.schurProbes;
            }

            if (m_borderInverse == nullptr) {
                m_borderInverse = std::make_unique<backend::PreparedDenseDirect>(m_borderBackend, m_schurComplement);
            } else {
                m_borderInverse->Refresh(m_schurComplement);
            }
            ++m_statistics.setups;
        }

        const mfem::Solver *m_structureInverse;
        const CouplingOperator *m_couplings;
        backend::DenseDirect m_borderBackend;
        mfem::DenseMatrix m_schurComplement;
        mfem::DenseMatrix m_structureInverseBorderCoupling;
        std::unique_ptr<backend::PreparedDenseDirect> m_borderInverse;
        mutable mfem::Vector m_structureWorkspace;
        mutable mfem::Vector m_structureCoupling;
        mutable mfem::Vector m_borderWorkspace;
        mutable mfem::Vector m_borderCoupling;
        mutable mfem::Vector m_borderDiagonal;
        mutable mfem::Vector m_borderBasis;
        mutable SpecificationBorderFactorizationStatistics m_statistics;
    };

    struct SpecificationBorderBlockPreparationReport final {
        bool structureRefreshed{false};
        bool rebuiltSchurComplement{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return structureRefreshed || rebuiltSchurComplement;
        }
    };

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

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem, SpecificationBorderBlockType Block>
    class PreparedSpecificationBorderBlock final : public mfem::Solver {
    private:
        using ProblemType       = std::remove_cvref_t<Problem>;
        using BlockType         = std::remove_cvref_t<Block>;
        using PreparedStructure = decltype(preconditioning::prepare(
            std::declval<const ProblemType &>(),
            std::declval<typename BlockType::StructureComponent>()
        ));
        static constexpr ApplicationContract structureInverseContract =
            backend::applicationContract<typename BlockType::StructureComponent::BackendType>;

    public:
        using Factorization = SpecificationBorderFactorizationOperator<
            SpecificationBorderJacobianOperator<ProblemType>,
            structureInverseContract>;

        PreparedSpecificationBorderBlock(
            const ProblemType &problem,
            BlockType block
        )
            : mfem::Solver(problem.StateSize()),
              m_problem(std::addressof(problem)),
              m_block(std::move(block)),
              m_structure(
                  preconditioning::prepare(
                      problem,
                      m_block.structureComponent()
                  )
              ),
              m_couplings(problem),
              m_factorization(
                  m_structure,
                  m_couplings,
                  m_block.borderBackend()
              ),
              m_snapshot(StellarEquilibriumProblemTraits<ProblemType>::Snapshot(problem)) {
            if (m_factorization.Height() != Height() || m_factorization.Width() != Width()) {
                throw std::logic_error(
                    "The prepared specification border does not span the grouped equilibrium coordinates."
                );
            }
            m_statistics.setups = 1;
        }

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

        [[nodiscard]] SpecificationBorderBlockPreparationReport Refresh() {
            const auto current = StellarEquilibriumProblemTraits<ProblemType>::Snapshot(*m_problem);
            ++m_statistics.refreshChecks;
            SpecificationBorderBlockPreparationReport report;
            const auto structureReport = m_structure.Refresh();
            report.structureRefreshed  = structureReport.DidAnyWork();

            if (current != m_snapshot) {
                m_factorization.RefreshSchurComplement();
                report.rebuiltSchurComplement = true;
                m_snapshot                    = current;
                ++m_statistics.refreshes;
            } else {
                ++m_statistics.noOpRefreshes;
            }
            return report;
        }

        [[nodiscard]] bool IsCurrent() const {
            return m_structure.IsCurrent() && m_problem->IsPrepared() &&
                   StellarEquilibriumProblemTraits<ProblemType>::Snapshot(*m_problem) == m_snapshot;
        }

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

        [[nodiscard]] const PreparedStructure &GetStructurePreconditioner() const noexcept {
            return m_structure;
        }

        [[nodiscard]] const SpecificationBorderJacobianOperator<ProblemType> &GetCouplings() const noexcept {
            return m_couplings;
        }

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

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

    private:
        const ProblemType *m_problem;
        BlockType m_block;
        PreparedStructure m_structure;
        SpecificationBorderJacobianOperator<ProblemType> m_couplings;
        Factorization m_factorization;
        StellarPreconditionerLifecycleSnapshot m_snapshot;
        PreparedSpecificationBorderBlockStatistics m_statistics;
    };

    template <
        equilibrium::DiscretizedStellarEquilibriumProblem Problem,
        PreconditionerComponent StructureComponent>
    [[nodiscard]] constexpr auto specificationBorderBlock(
        const Problem &,
        StructureComponent structureComponent,
        backend::DenseDirect borderBackend = {}
    ) {
        using ProblemType = std::remove_cvref_t<Problem>;
        using Block       = SpecificationBorderBlock<
            StructureComponent, typename ProblemType::ModelType, typename ProblemType::FormType,
            typename ProblemType::JacobianFormType>;
        using Plan = PreconditionerPlan<Block>;
        static_assert(
            CompletePreconditionerFor<Plan, typename ProblemType::FormType>,
            "The model-compiled preconditioner must own every correction and residual block exactly once."
        );
        static_assert(
            CompatiblePreconditionerFor<Plan, typename ProblemType::FormType, typename ProblemType::JacobianFormType>,
            "Every coupling required by the model-compiled preconditioner must exist in the compiled Jacobian."
        );
        return Block{std::move(structureComponent), std::move(borderBackend)};
    }

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
    [[nodiscard]] constexpr auto specificationBorderBlock(const Problem &problem) {
        return specificationBorderBlock(problem, stellarStructureBlock(problem), backend::DenseDirect{});
    }

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
    [[nodiscard]] constexpr auto makePreconditioner(const Problem &problem) {
        return specificationBorderBlock(problem);
    }

} // namespace mean_field::preconditioning