module;

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

#include <mfem.hpp>

export module mean_field:normalization.stellar_equilibrium;

export import :normalization.operators;
export import :operators.stellar_equilibrium_compiler;
export import :operators.stellar_equilibrium_problem;
export import :utils.domain;

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

    [[nodiscard]] inline mfem::Vector AssembleScalarMassDiagonal(
        mfem::ParFiniteElementSpace &space,
        mfem::Array<int> *domainMarker = nullptr
    ) {
        mfem::ParBilinearForm mass(&space);
        if (domainMarker == nullptr) {
            mass.AddDomainIntegrator(new mfem::MassIntegrator());
        } else {
            mass.AddDomainIntegrator(new mfem::MassIntegrator(), *domainMarker);
        }
        mass.Assemble();
        mass.Finalize();
        std::unique_ptr<mfem::HypreParMatrix> matrix(mass.ParallelAssemble());
        if (matrix == nullptr) {
            throw std::runtime_error("Reference scalar Riesz mass assembly failed.");
        }
        mfem::Vector diagonal;
        matrix->GetDiag(diagonal);
        return diagonal;
    }

    [[nodiscard]] inline mfem::Vector AssembleHDivMassDiagonal(mfem::ParFiniteElementSpace &space) {
        mfem::ParBilinearForm mass(&space);
        mass.AddDomainIntegrator(new mfem::VectorFEMassIntegrator());
        mass.Assemble();
        mass.Finalize();
        std::unique_ptr<mfem::HypreParMatrix> matrix(mass.ParallelAssemble());
        if (matrix == nullptr) {
            throw std::runtime_error("Reference H(div) Riesz mass assembly failed.");
        }
        mfem::Vector diagonal;
        matrix->GetDiag(diagonal);
        return diagonal;
    }

    [[nodiscard]] inline mfem::Vector AssembleSurfaceMassDiagonal(
        const fem::FEM &finiteElements,
        const field::ScalarBoundaryDofMap &surfaceMap
    ) {
        mfem::Array<int> marker(finiteElements.mesh->bdr_attributes.Max());
        marker = 0;
        constexpr int attribute = DomainSchema::template boundary_attribute<utils::domain::StellarSurface>();
        if (attribute <= 0 || attribute > marker.Size()) {
            throw std::invalid_argument("The reference mesh does not contain the stellar-surface boundary.");
        }
        marker[attribute - 1] = 1;

        mfem::ParBilinearForm mass(finiteElements.surfaceDeformationFes.get());
        mass.AddBoundaryIntegrator(new mfem::MassIntegrator(), marker);
        mass.Assemble();
        mass.Finalize();
        std::unique_ptr<mfem::HypreParMatrix> matrix(mass.ParallelAssemble());
        if (matrix == nullptr) {
            throw std::runtime_error("Reference surface Riesz mass assembly failed.");
        }
        mfem::Vector ambientDiagonal;
        matrix->GetDiag(ambientDiagonal);
        return surfaceMap.gather(ambientDiagonal);
    }

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
    [[nodiscard]] const auto &PhysicalOperator(const Problem &problem) {
        return problem.GetPhysicalOperator();
    }

    [[nodiscard]] inline mfem::Vector GatherDiagonal(
        const mfem::Vector &fullDiagonal,
        const field::FieldDofMap &map,
        const char *role
    ) {
        if (fullDiagonal.Size() != map.full_size()) {
            throw std::logic_error(std::string("The reference ") + role + " Gram diagonal has an incompatible map.");
        }
        return map.gather(fullDiagonal);
    }
} // namespace mean_field::normalization::detail

export namespace mean_field::normalization {
    /*
     * Runtime preparation paired with the compile-time normalization plan.
     * The operator compiler is the authority for which blocks a specification
     * generated, and PhysicalRieszBlockTraits is the authority for their
     * declared physical laws.  Keeping those responsibilities separate means
     * this layer never names a concrete integral or phase constraint.
     */
    namespace detail {
        template <typename Block>
        using PhysicalRieszMethodFor = typename PhysicalRieszBlockTraits<Block>::Method;

        template <typename Block, typename = void>
        struct IsGlobalGeneratedValueNormalization : std::false_type { };

        template <typename Generated>
        struct IsGlobalGeneratedValueNormalization<
            utils::blocks::generated_value_block<Generated>,
            std::void_t<
                decltype(PhysicalRieszMethodFor<
                         utils::blocks::generated_value_block<Generated>>::topology),
                decltype(PhysicalRieszMethodFor<
                         utils::blocks::generated_value_block<Generated>>::scale)>>
            : std::bool_constant<
                  PhysicalRieszBlockTraits<
                      utils::blocks::generated_value_block<Generated>>::registered &&
                  PhysicalRieszMethodFor<
                      utils::blocks::generated_value_block<Generated>>::topology ==
                      RieszTopology::global_scalar> { };

        template <typename Blocks, typename Specification>
        struct GeneratedValueBlocksBelongToSpecification : std::false_type { };

        template <typename Generated, typename Specification, typename = void>
        struct GeneratedCoordinateBelongsToSpecification : std::false_type { };

        template <typename Generated, typename Specification>
        struct GeneratedCoordinateBelongsToSpecification<
            Generated,
            Specification,
            std::void_t<typename Generated::SpecificationType>>
            : std::bool_constant<
                  std::same_as<typename Generated::SpecificationType, Specification>> { };

        template <typename Specification, typename... Generated>
        struct GeneratedValueBlocksBelongToSpecification<
            utils::blocks::type_list<utils::blocks::generated_value_block<Generated>...>,
            Specification>
            : std::bool_constant<
                  (GeneratedCoordinateBelongsToSpecification<Generated, Specification>::value && ...)> { };

        template <typename Block, typename = void>
        struct IsGlobalGeneratedResidualNormalization : std::false_type { };

        template <typename Generated>
        struct IsGlobalGeneratedResidualNormalization<
            utils::blocks::generated_residual_block<Generated>,
            std::void_t<
                decltype(PhysicalRieszMethodFor<
                         utils::blocks::generated_residual_block<Generated>>::topology),
                decltype(PhysicalRieszMethodFor<
                         utils::blocks::generated_residual_block<Generated>>::scale)>>
            : std::bool_constant<
                  PhysicalRieszBlockTraits<
                      utils::blocks::generated_residual_block<Generated>>::registered &&
                  PhysicalRieszMethodFor<
                      utils::blocks::generated_residual_block<Generated>>::topology ==
                      RieszTopology::global_scalar> { };

        template <typename Blocks, typename Specification>
        struct GeneratedResidualBlocksBelongToSpecification : std::false_type { };

        template <typename Specification, typename... Generated>
        struct GeneratedResidualBlocksBelongToSpecification<
            utils::blocks::type_list<utils::blocks::generated_residual_block<Generated>...>,
            Specification>
            : std::bool_constant<
                  (GeneratedCoordinateBelongsToSpecification<Generated, Specification>::value && ...)> { };

        template <typename Blocks> struct PrepareGeneratedValueNormalizations {
            static constexpr bool registered = false;

            template <typename Form>
            static constexpr bool completeFor = false;

            template <typename Form>
            static void Apply(
                DiagonalNormalizationBuilder<Form> &,
                const StellarCharacteristicScales &
            ) {
                static_assert(registered, "Generated value-block normalization metadata is malformed.");
            }
        };

        template <typename... Blocks>
        struct PrepareGeneratedValueNormalizations<utils::blocks::type_list<Blocks...>> {
            static constexpr bool registered =
                (IsGlobalGeneratedValueNormalization<Blocks>::value && ...);

            template <typename Form>
            static constexpr bool completeFor = registered &&
                utils::blocks::block_form_is_valid_v<Form> &&
                (utils::blocks::contains_type_v<Blocks, typename Form::value_blocks> && ...);

            template <typename Form>
            static void Apply(
                DiagonalNormalizationBuilder<Form> &builder,
                const StellarCharacteristicScales &scales
            ) {
                if constexpr (completeFor<Form>) {
                    (builder.template SetValueGlobal<Blocks>(physicalScale<Blocks>(scales)), ...);
                } else {
                    static_assert(
                        completeFor<Form>,
                        "Every generated value block must have a declared global-scalar Physical Riesz law "
                        "and belong to the compiled equilibrium form."
                    );
                }
            }
        };

        template <typename Blocks> struct PrepareGeneratedResidualNormalizations {
            static constexpr bool registered = false;

            template <typename Form>
            static constexpr bool completeFor = false;

            template <typename Form>
            static void Apply(
                DiagonalNormalizationBuilder<Form> &,
                const StellarCharacteristicScales &
            ) {
                static_assert(registered, "Generated residual-block normalization metadata is malformed.");
            }
        };

        template <typename... Blocks>
        struct PrepareGeneratedResidualNormalizations<utils::blocks::type_list<Blocks...>> {
            static constexpr bool registered =
                (IsGlobalGeneratedResidualNormalization<Blocks>::value && ...);

            template <typename Form>
            static constexpr bool completeFor = registered &&
                utils::blocks::block_form_is_valid_v<Form> &&
                (utils::blocks::contains_type_v<Blocks, typename Form::residual_blocks> && ...);

            template <typename Form>
            static void Apply(
                DiagonalNormalizationBuilder<Form> &builder,
                const StellarCharacteristicScales &scales
            ) {
                if constexpr (completeFor<Form>) {
                    (builder.template SetResidualGlobal<Blocks>(physicalScale<Blocks>(scales)), ...);
                } else {
                    static_assert(
                        completeFor<Form>,
                        "Every generated residual block must have a declared global-scalar Physical Riesz law "
                        "and belong to the compiled equilibrium form."
                    );
                }
            }
        };

        template <typename Specification, typename = void>
        struct CompileStellarSpecificationNormalization {
            using ValuePreparation = PrepareGeneratedValueNormalizations<void>;
            using ResidualPreparation = PrepareGeneratedResidualNormalizations<void>;

            static constexpr bool registered = false;

            template <typename Form>
            static constexpr bool completeFor = false;

            template <typename Form>
            static void Apply(
                DiagonalNormalizationBuilder<Form> &,
                const StellarCharacteristicScales &
            ) {
                static_assert(
                    completeFor<Form>,
                    "The specification has no complete generated-coordinate normalization."
                );
            }
        };

        template <models::ModelSpecification Specification>
        struct CompileStellarSpecificationNormalization<
            Specification,
            std::void_t<
                typename operators::StellarEquilibriumSpecificationCompilation<
                    Specification>::GeneratedValueBlocks,
                typename operators::StellarEquilibriumSpecificationCompilation<
                    Specification>::GeneratedResidualBlocks>> {
            using OperatorCompilation =
                operators::StellarEquilibriumSpecificationCompilation<Specification>;
            using ValuePreparation = PrepareGeneratedValueNormalizations<
                typename OperatorCompilation::GeneratedValueBlocks>;
            using ResidualPreparation = PrepareGeneratedResidualNormalizations<
                typename OperatorCompilation::GeneratedResidualBlocks>;

            static constexpr bool registered = OperatorCompilation::complete &&
                                               models::CompleteGeneratedNormalizationFor<
                                                   Specification> &&
                                               GeneratedValueBlocksBelongToSpecification<
                                                   typename OperatorCompilation::GeneratedValueBlocks,
                                                   Specification>::value &&
                                               GeneratedResidualBlocksBelongToSpecification<
                                                   typename OperatorCompilation::GeneratedResidualBlocks,
                                                   Specification>::value &&
                                               ValuePreparation::registered &&
                                               ResidualPreparation::registered;

            template <typename Form>
            static constexpr bool completeFor = registered &&
                                                ValuePreparation::template completeFor<Form> &&
                                                ResidualPreparation::template completeFor<Form>;

            template <typename Form>
            static void Apply(
                DiagonalNormalizationBuilder<Form> &builder,
                const StellarCharacteristicScales &scales
            ) {
                if constexpr (completeFor<Form>) {
                    ValuePreparation::template Apply<Form>(builder, scales);
                    ResidualPreparation::template Apply<Form>(builder, scales);
                } else {
                    static_assert(
                        completeFor<Form>,
                        "The specification's generated blocks do not have a complete runtime normalization."
                    );
                }
            }
        };

        template <typename SpecificationSet> struct PrepareSpecificationNormalizations;

        template <models::ModelSpecification... Specifications>
        struct PrepareSpecificationNormalizations<models::detail::SpecificationSetStorage<Specifications...>> {
            static constexpr bool registered =
                (CompileStellarSpecificationNormalization<Specifications>::registered && ...);

            template <typename Form>
            static constexpr bool completeFor =
                (CompileStellarSpecificationNormalization<Specifications>::template completeFor<Form> && ...);

            template <typename Form>
            static void Apply(
                DiagonalNormalizationBuilder<Form> &builder,
                const StellarCharacteristicScales &scales
            ) {
                static_assert(
                    completeFor<Form>,
                    "Every generated stellar-equilibrium coordinate requires a declared global-scalar "
                    "Physical Riesz normalization and compiler-owned root block."
                );
                (CompileStellarSpecificationNormalization<Specifications>::template Apply<Form>(builder, scales), ...);
            }
        };

        template <typename Model, typename Form, typename = void>
        struct StellarModelNormalizationCoverage : std::false_type { };

        template <typename Model, typename Form>
            requires model::StellarModelType<Model> && utils::blocks::block_form_is_valid_v<Form>
        struct StellarModelNormalizationCoverage<
            Model,
            Form,
            std::void_t<typename std::remove_cvref_t<Model>::SpecificationTypes>>
            : std::bool_constant<
                  PrepareSpecificationNormalizations<
                      typename std::remove_cvref_t<Model>::SpecificationTypes>::template completeFor<Form>> { };
    } // namespace detail

    template <typename Specification>
    struct StellarSpecificationNormalizationContribution
        : detail::CompileStellarSpecificationNormalization<std::remove_cvref_t<Specification>> {
        using Base = detail::CompileStellarSpecificationNormalization<std::remove_cvref_t<Specification>>;

        template <typename Form>
        static void Apply(
            DiagonalNormalizationBuilder<Form> &builder,
            const StellarCharacteristicScales &scales
        ) {
            static_assert(
                Base::template completeFor<Form>,
                "The specification's generated blocks do not have a complete runtime normalization."
            );
            Base::template Apply<Form>(builder, scales);
        }
    };

    template <typename Specification>
    concept RegisteredStellarSpecificationNormalization =
        StellarSpecificationNormalizationContribution<Specification>::registered;

    template <typename Specification, typename Form>
    concept CompleteStellarSpecificationNormalizationFor =
        utils::blocks::block_form_is_valid_v<Form> &&
        StellarSpecificationNormalizationContribution<Specification>::template completeFor<Form>;

    template <typename Model, typename Form>
    concept CompleteStellarNormalizationFor =
        detail::StellarModelNormalizationCoverage<
            std::remove_cvref_t<Model>,
            std::remove_cvref_t<Form>>::value;

    /*
     * Physical Riesz preparation is an optional capability of a physical
     * core, not part of the protocol needed by the variadic equilibrium root.
     * Keeping this boundary structural lets a new EOS core opt in by exposing
     * the same discretization maps without inheriting from, or otherwise
     * naming, the Polytrope implementation.
     */
    template <typename Candidate>
    concept PhysicalRieszStellarEquilibriumCore =
        operators::PreparedStellarEquilibriumPhysicalCore<std::remove_cvref_t<Candidate>> &&
        PhysicalRieszCoreRuntime<std::remove_cvref_t<Candidate>>;

    template <typename Problem>
    concept PhysicalRieszStellarEquilibriumProblem =
        equilibrium::DiscretizedStellarEquilibriumProblem<std::remove_cvref_t<Problem>> &&
        requires {
            typename std::remove_cvref_t<Problem>::ModelType;
            typename std::remove_cvref_t<Problem>::FormType;
            typename std::remove_cvref_t<Problem>::PhysicalCoreType;
            typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType;
            requires PhysicalRieszDiagonalPrescription<
                typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType>;
            requires CompilableNormalizationFor<
                typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType,
                typename std::remove_cvref_t<Problem>::FormType>;
            requires CompleteStellarNormalizationFor<
                typename std::remove_cvref_t<Problem>::ModelType,
                typename std::remove_cvref_t<Problem>::FormType>;
            requires StellarNormalizationRuntimeAvailableFor<
                typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType,
                typename std::remove_cvref_t<Problem>::FormType,
                typename std::remove_cvref_t<Problem>::PhysicalCoreType,
                typename std::remove_cvref_t<Problem>::ModelType::SpecificationTypes>;
        };

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        requires std::same_as<
            typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType,
            Unnormalized>
    [[nodiscard]] DiagonalNormalization prepareNormalization(const Problem &problem) {
        return DiagonalNormalization::Identity(problem.StateSize(), problem.EquationSize());
    }

    template <PhysicalRieszStellarEquilibriumProblem Problem>
    [[nodiscard]] DiagonalNormalization prepareNormalization(const Problem &problem) {
        using ProblemType = std::remove_cvref_t<Problem>;
        using Form = typename ProblemType::FormType;

        const fem::FEM &finiteElements = problem.GetDiscretization().finiteElementModel();
        if (!finiteElements.okay()) {
            throw std::invalid_argument("Physical Riesz preparation requires a current finite-element model.");
        }

        const auto &physical = detail::PhysicalOperator(problem);
        const auto &gravityContext = physical.GetGravityContext();
        const auto &enthalpyMap = physical.GetHydrostaticOperator().GetEnthalpyMap();
        const auto scales = deriveStellarCharacteristicScales(
            problem.GetNormalizationPrescription(),
            problem.GetStellarModel()
        );

        mfem::Array<int> stellarMarker =
            utils::domain::make_attribute_marker<utils::domain::Stellar, detail::DomainSchema>(*finiteElements.mesh);
        const mfem::Vector densityDiagonal = detail::GatherDiagonal(
            detail::AssembleScalarMassDiagonal(*finiteElements.densityFes, &stellarMarker),
            gravityContext.GetDensityMap(),
            "density"
        );
        const mfem::Vector enthalpyDiagonal = detail::GatherDiagonal(
            detail::AssembleScalarMassDiagonal(*finiteElements.enthalpyFes, &stellarMarker),
            enthalpyMap,
            "enthalpy"
        );
        const mfem::Vector gravityGradientDiagonal = detail::GatherDiagonal(
            detail::AssembleHDivMassDiagonal(*finiteElements.gravityFluxFes),
            gravityContext.GetGravityGradientMap(),
            "gravity-gradient"
        );
        const mfem::Vector gravityPotentialDiagonal = detail::GatherDiagonal(
            detail::AssembleScalarMassDiagonal(*finiteElements.gravityPotentialFes),
            gravityContext.GetGravityPotentialMap(),
            "gravity-potential"
        );
        const field::ScalarBoundaryDofMap surfaceMap =
            field::make_stellar_surface_scalar_dof_map<detail::DomainSchema>(*finiteElements.surfaceDeformationFes);
        const mfem::Vector surfaceDiagonal = detail::AssembleSurfaceMassDiagonal(finiteElements, surfaceMap);
        if (surfaceDiagonal.Size() != physical.GetDomainDeformation().parameterCount()) {
            throw std::logic_error("The reference surface Gram diagonal does not match the root surface block.");
        }

        DiagonalNormalizationBuilder<Form> builder(problem.GetManifest().layout());
        builder.template SetValueBlock<utils::blocks::density::mass::value>(
            physicalScale<utils::blocks::density::mass::value>(scales), densityDiagonal
        );
        builder.template SetValueBlock<utils::blocks::surface_deformation::parameters::value>(
            physicalScale<utils::blocks::surface_deformation::parameters::value>(scales), surfaceDiagonal
        );
        builder.template SetValueBlock<utils::blocks::gravity::gradient::value>(
            physicalScale<utils::blocks::gravity::gradient::value>(scales), gravityGradientDiagonal
        );
        builder.template SetValueBlock<utils::blocks::gravity::poisson::value>(
            physicalScale<utils::blocks::gravity::poisson::value>(scales), gravityPotentialDiagonal
        );
        builder.template SetValueBlock<utils::blocks::enthalpy::specific::value>(
            physicalScale<utils::blocks::enthalpy::specific::value>(scales), enthalpyDiagonal
        );
        builder.template SetResidualBlock<utils::blocks::gravity::gradient::residual>(
            physicalScale<utils::blocks::gravity::gradient::residual>(scales), gravityGradientDiagonal
        );
        builder.template SetResidualBlock<utils::blocks::gravity::poisson::residual>(
            physicalScale<utils::blocks::gravity::poisson::residual>(scales), gravityPotentialDiagonal
        );
        builder.template SetResidualBlock<utils::blocks::density::mass::residual>(
            physicalScale<utils::blocks::density::mass::residual>(scales), densityDiagonal
        );
        builder.template SetResidualBlock<utils::blocks::surface_deformation::shape_equilibrium::residual>(
            physicalScale<utils::blocks::surface_deformation::shape_equilibrium::residual>(scales), surfaceDiagonal
        );
        const mfem::Array<int> &surfaceRows = problem.GetPressureSurfaceRows().reduced_dofs();
        builder.template SetHybridResidualBlock<utils::blocks::enthalpy::specific::residual>(
            physicalScale<utils::blocks::enthalpy::specific::residual>(scales),
            enthalpyDiagonal,
            std::span<const int>{surfaceRows.GetData(), static_cast<std::size_t>(surfaceRows.Size())}
        );
        detail::PrepareSpecificationNormalizations<typename ProblemType::ModelType::SpecificationTypes>::Apply(
            builder,
            scales
        );

        return std::move(builder).Build();
    }

    /* Public adapter for a third-party prescription. The implementation stays
     * beside the policy and has the readable signature
     *
     *   prepareStellarNormalization(policy, problem)
     *
     * while every solver-facing caller continues to use the uniform
     * prepareNormalization(problem) operation. */
    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        requires(
            !std::same_as<
                typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType,
                Unnormalized> &&
            !PhysicalRieszDiagonalPrescription<
                typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType> &&
            RuntimePreparedNormalizationOperation<Problem>)
    [[nodiscard]] DiagonalNormalization prepareNormalization(
        const Problem &problem
    ) {
        return prepareStellarNormalization(
            problem.GetNormalizationPrescription(),
            problem
        );
    }

    /*
     * Solver-facing normalization exists exactly when runtime preparation for
     * the problem's compile-time prescription is a valid operation.  This
     * folds future policy-owned preparation hooks into the same public contract and
     * turns unsupported core/prescription pairs into ordinary constraint
     * failure instead of an error in a constructor body.
     */
    template <typename Problem>
    concept NormalizableStellarEquilibriumProblem =
        equilibrium::DiscretizedStellarEquilibriumProblem<std::remove_cvref_t<Problem>> &&
        requires(const std::remove_cvref_t<Problem> &problem) {
            {
                prepareNormalization(problem)
            } -> std::same_as<DiagonalNormalization>;
        };

    struct NormalizedStellarEquilibriumStatistics final {
        std::uint64_t normalizationPreparations{0};
        std::uint64_t physicalPreparations{0};
        std::uint64_t residualRetrievals{0};
        std::uint64_t jacobianApplications{0};
    };

    /*
     * The high-level stellar adapter retains a pointer to a prepared inverse.
     * Consequently that inverse must identify the exact physical problem and
     * expose its lifecycle state.  Generic MFEM solvers remain valid inputs to
     * the lower-level ScaledPreconditioner, where no stellar association is
     * implied.
     */
    template <typename Candidate, typename Problem>
    concept ProblemBoundStellarInverseFor =
        NormalizableStellarEquilibriumProblem<std::remove_cvref_t<Problem>> &&
        std::derived_from<std::remove_cvref_t<Candidate>, mfem::Solver> &&
        requires(const std::remove_cvref_t<Candidate> &inverse) {
            {
                inverse.GetProblem()
            } -> std::same_as<const std::remove_cvref_t<Problem> &>;
            {
                inverse.IsCurrent()
            } -> std::same_as<bool>;
        };

    template <NormalizableStellarEquilibriumProblem Problem, typename PhysicalInverse>
        requires ProblemBoundStellarInverseFor<PhysicalInverse, Problem>
    class NormalizedStellarPreconditioner;

    /*
     * Solver-facing coordinates for a dimensional stellar problem. The
     * physical problem remains the sole source of residual and Jacobian
     * physics; this adapter performs only the coordinate maps
     *
     *     x = R x_hat,   F_hat = L F,   J_hat = L J R.
     *
     * Its normalization is immutable during Prepare/BuildResidual/Mult and is
     * changed only by an explicit RefreshNormalization call.
     */
    template <NormalizableStellarEquilibriumProblem Problem>
    class NormalizedStellarEquilibriumOperator final : public mfem::Operator {
    private:
        using ProblemType = std::remove_cvref_t<Problem>;

    public:
        explicit NormalizedStellarEquilibriumOperator(ProblemType &problem)
            : mfem::Operator(problem.EquationSize(), problem.StateSize()),
              m_problem(&problem),
              m_normalization(prepareNormalization(problem)),
              m_scaledJacobian(problem.GetLinearizationOperator(), m_normalization),
              m_physicalState(problem.StateSize()),
              m_physicalResidual(problem.EquationSize()),
              m_normalizedResidual(problem.EquationSize()) {
            if (Width() != Height()) {
                throw std::invalid_argument("A normalized stellar-equilibrium operator must be square.");
            }
            m_statistics.normalizationPreparations = 1;
        }

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

        [[nodiscard]] auto Prepare(
            const mfem::Vector &normalizedState,
            const operators::StellarEquilibriumDependencies &dependencies,
            const physics::RigidRotation &rotation
        ) requires(ProblemType::generatedRotationProviderCount == 0) {
            if (normalizedState.Size() != Width()) {
                throw std::invalid_argument("The normalized stellar state has the wrong size.");
            }

            m_isPrepared = false;
            m_normalization.DenormalizeState(normalizedState, m_physicalState);
            auto report = m_problem->Prepare(m_physicalState, dependencies, rotation);
            m_problem->BuildResidual(m_physicalResidual);
            m_normalization.NormalizeResidual(m_physicalResidual, m_normalizedResidual);
            m_physicalPreparationGeneration = m_problem->GetPreparationGeneration();
            m_isPrepared                    = true;
            ++m_statistics.physicalPreparations;
            return report;
        }

        [[nodiscard]] auto Prepare(
            const mfem::Vector &normalizedState,
            const operators::StellarEquilibriumDependencies &dependencies
        ) requires(ProblemType::generatedRotationProviderCount == 1) {
            if (normalizedState.Size() != Width()) {
                throw std::invalid_argument("The normalized stellar state has the wrong size.");
            }

            m_isPrepared = false;
            m_normalization.DenormalizeState(normalizedState, m_physicalState);
            auto report = m_problem->Prepare(m_physicalState, dependencies);
            m_problem->BuildResidual(m_physicalResidual);
            m_normalization.NormalizeResidual(m_physicalResidual, m_normalizedResidual);
            m_physicalPreparationGeneration = m_problem->GetPreparationGeneration();
            m_isPrepared                    = true;
            ++m_statistics.physicalPreparations;
            return report;
        }

        void BuildResidual(mfem::Vector &normalizedResidual) const {
            VerifyPrepared();
            normalizedResidual = m_normalizedResidual;
            ++m_statistics.residualRetrievals;
        }

        void Mult(
            const mfem::Vector &normalizedDirection,
            mfem::Vector &normalizedAction
        ) const override {
            VerifyPrepared();
            if (normalizedDirection.Size() != Width()) {
                throw std::invalid_argument("The normalized stellar direction has the wrong size.");
            }
            m_scaledJacobian.Mult(normalizedDirection, normalizedAction);
            ++m_statistics.jacobianApplications;
        }

        void RefreshNormalization() {
            DiagonalNormalization refreshed = prepareNormalization(*m_problem);
            m_normalization                  = std::move(refreshed);
            m_isPrepared                     = false;
            ++m_statistics.normalizationPreparations;
        }

        void NormalizeState(
            const mfem::Vector &physicalState,
            mfem::Vector &normalizedState
        ) const {
            m_normalization.NormalizeState(physicalState, normalizedState);
        }

        void DenormalizeState(
            const mfem::Vector &normalizedState,
            mfem::Vector &physicalState
        ) const {
            m_normalization.DenormalizeState(normalizedState, physicalState);
        }

        void NormalizeResidual(
            const mfem::Vector &physicalResidual,
            mfem::Vector &normalizedResidual
        ) const {
            m_normalization.NormalizeResidual(physicalResidual, normalizedResidual);
        }

        void DenormalizeResidual(
            const mfem::Vector &normalizedResidual,
            mfem::Vector &physicalResidual
        ) const {
            m_normalization.DenormalizeResidual(normalizedResidual, physicalResidual);
        }

        template <typename PhysicalInverse>
            requires ProblemBoundStellarInverseFor<PhysicalInverse, Problem>
        [[nodiscard]] NormalizedStellarPreconditioner<Problem, std::remove_cvref_t<PhysicalInverse>>
        MakeScaledPreconditioner(PhysicalInverse &physicalInverse) const;

        [[nodiscard]] bool IsPrepared() const noexcept {
            return m_isPrepared && m_problem->IsPrepared() &&
                   m_physicalPreparationGeneration == m_problem->GetPreparationGeneration();
        }

        [[nodiscard]] ProblemType &GetPhysicalProblem() noexcept {
            return *m_problem;
        }

        [[nodiscard]] const ProblemType &GetPhysicalProblem() const noexcept {
            return *m_problem;
        }

        [[nodiscard]] const mfem::Operator &GetPhysicalJacobian() const noexcept {
            return m_problem->GetLinearizationOperator();
        }

        [[nodiscard]] const ProblemType &GetProblem() const noexcept {
            return *m_problem;
        }

        [[nodiscard]] const DiagonalNormalization &GetNormalization() const noexcept {
            return m_normalization;
        }

        [[nodiscard]] const mfem::Vector &GetPhysicalState() const {
            VerifyPrepared();
            return m_physicalState;
        }

        [[nodiscard]] const mfem::Vector &GetPhysicalResidual() const {
            VerifyPrepared();
            return m_physicalResidual;
        }

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

    private:
        void VerifyPrepared() const {
            if (!IsPrepared()) {
                throw std::logic_error(
                    "The normalized stellar-equilibrium operator must be prepared and current before application."
                );
            }
        }

        ProblemType *m_problem;
        DiagonalNormalization m_normalization;
        ScaledJacobianOperator m_scaledJacobian;
        mfem::Vector m_physicalState;
        mfem::Vector m_physicalResidual;
        mfem::Vector m_normalizedResidual;
        std::uint64_t m_physicalPreparationGeneration{0};
        mutable NormalizedStellarEquilibriumStatistics m_statistics;
        bool m_isPrepared{false};
    };

    template <NormalizableStellarEquilibriumProblem Problem, typename PhysicalInverse>
        requires ProblemBoundStellarInverseFor<PhysicalInverse, Problem>
    class NormalizedStellarPreconditioner final : public mfem::Solver {
    private:
        using ProblemType = std::remove_cvref_t<Problem>;
        using NormalizedOperator = NormalizedStellarEquilibriumOperator<ProblemType>;
        using PhysicalInverseType = std::remove_cvref_t<PhysicalInverse>;

        [[nodiscard]] static PhysicalInverseType &RequireAssociatedPhysicalInverse(
            const NormalizedOperator &normalizedOperator,
            PhysicalInverseType &physicalInverse
        ) {
            if (std::addressof(physicalInverse.GetProblem()) !=
                std::addressof(normalizedOperator.GetProblem())) {
                throw std::invalid_argument(
                    "A normalized stellar preconditioner and its physical inverse must belong to the same problem."
                );
            }
            return physicalInverse;
        }

    public:
        NormalizedStellarPreconditioner(
            const NormalizedOperator &normalizedOperator,
            PhysicalInverseType &physicalInverse
        )
            : mfem::Solver(
                  normalizedOperator.Width(),
                  normalizedOperator.Height(),
                  physicalInverse.iterative_mode
              ),
              m_normalizedOperator(&normalizedOperator),
              m_physicalInverse(&physicalInverse),
              m_scaled(
                  RequireAssociatedPhysicalInverse(normalizedOperator, physicalInverse),
                  normalizedOperator.GetPhysicalJacobian(),
                  normalizedOperator,
                  normalizedOperator.GetNormalization()
              ) {
        }

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

        void SetOperator(const mfem::Operator &normalizedJacobian) override {
            VerifyCurrent();
            if (&normalizedJacobian != m_normalizedOperator) {
                throw std::invalid_argument(
                    "The normalized stellar preconditioner cannot be rebound to a different Jacobian."
                );
            }
            m_scaled.SetOperator(normalizedJacobian);
        }

        void Mult(
            const mfem::Vector &normalizedResidual,
            mfem::Vector &normalizedCorrection
        ) const override {
            VerifyCurrent();
            m_scaled.Mult(normalizedResidual, normalizedCorrection);
        }

        [[nodiscard]] bool IsCurrent() const {
            return m_normalizedOperator->IsPrepared() &&
                   m_physicalInverse->IsCurrent();
        }

        [[nodiscard]] PhysicalInverseType &GetPhysicalInverse() noexcept {
            return *m_physicalInverse;
        }

        [[nodiscard]] const PhysicalInverseType &GetPhysicalInverse() const noexcept {
            return *m_physicalInverse;
        }

        [[nodiscard]] const mfem::Operator &GetPhysicalJacobian() const noexcept {
            return m_scaled.GetPhysicalJacobian();
        }

        [[nodiscard]] const mfem::Operator &GetNormalizedJacobian() const {
            return m_scaled.GetNormalizedJacobian();
        }

        [[nodiscard]] const ScaledPreconditionerStatistics &GetStatistics() const noexcept {
            return m_scaled.GetStatistics();
        }

    private:
        void VerifyCurrent() const {
            if (!IsCurrent()) {
                throw std::logic_error(
                    "The normalized stellar preconditioner cannot be used while its normalized operator or physical "
                    "inverse is stale."
                );
            }
        }

        const NormalizedOperator *m_normalizedOperator;
        PhysicalInverseType *m_physicalInverse;
        ScaledPreconditioner m_scaled;
    };

    template <NormalizableStellarEquilibriumProblem Problem>
    template <typename PhysicalInverse>
        requires ProblemBoundStellarInverseFor<PhysicalInverse, Problem>
    NormalizedStellarPreconditioner<Problem, std::remove_cvref_t<PhysicalInverse>>
    NormalizedStellarEquilibriumOperator<Problem>::MakeScaledPreconditioner(PhysicalInverse &physicalInverse) const {
        VerifyPrepared();
        return NormalizedStellarPreconditioner<Problem, std::remove_cvref_t<PhysicalInverse>>{
            *this,
            physicalInverse
        };
    }

    template <NormalizableStellarEquilibriumProblem Problem>
    [[nodiscard]] auto makeNormalizedStellarEquilibriumOperator(Problem &problem) {
        return NormalizedStellarEquilibriumOperator<Problem>{problem};
    }
} // namespace mean_field::normalization