module;

#include <array>
#include <cmath>
#include <cstdint>
#include <expected>
#include <stdexcept>
#include <utility>

#include <mfem.hpp>

module mean_field;

import :operators.prepared_stellar_equilibrium;

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

    void verify_coupled_discretization(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.mesh != nullptr && f.densityFes != nullptr && f.displacementFes != nullptr &&
                f.gravityFluxFes != nullptr && f.gravityPotentialFes != nullptr && f.enthalpyFes != nullptr,
            "PreparedStellarEquilibriumOperator requires the complete coupled finite-element discretization."
        );
    }

    using StellarRootForm = mean_field::utils::blocks::surface_deformed_stellar_equilibrium_form;

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection with_preparation_stage(
        mean_field::operators::StellarEquilibriumPreparationRejection rejection,
        const mean_field::operators::StellarEquilibriumPreparationStage stage
    ) noexcept {
        rejection.stage = stage;
        return rejection;
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection
    make_thermodynamic_rejection(const mean_field::eos::EvaluationErrorCode code) {
        using Failure = mean_field::operators::StellarEquilibriumPreparationRejection;
        using Reason  = mean_field::operators::StellarEquilibriumPreparationRejectionReason;

        switch (code) {
        case mean_field::eos::EvaluationErrorCode::outside_domain:
            return Failure{.reason = Reason::thermodynamic_domain, .thermodynamicErrorCode = code};
        case mean_field::eos::EvaluationErrorCode::nonfinite_input:
        case mean_field::eos::EvaluationErrorCode::nonfinite_result:
            return Failure{.reason = Reason::non_finite_thermodynamics, .thermodynamicErrorCode = code};
        default:
            throw std::logic_error(
                "A non-retryable equation-of-state error was incorrectly returned as a stellar trial rejection."
            );
        }
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection
    make_mapping_rejection(const mean_field::mapping::MappingStatus status) {
        using Failure = mean_field::operators::StellarEquilibriumPreparationRejection;
        using Reason  = mean_field::operators::StellarEquilibriumPreparationRejectionReason;
        return Failure{
            .reason = status == mean_field::mapping::MappingStatus::non_positive_determinant
                          ? Reason::inverted_geometry
                          : Reason::non_finite_geometry
        };
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection make_gravity_rejection(
        const mean_field::operators::context::gravity_field::GravityFieldPreparationRejection &rejection
    ) {
        using ChildReason = mean_field::operators::context::gravity_field::GravityFieldPreparationRejectionReason;
        if (rejection.reason == ChildReason::invalid_mapping) {
            return make_mapping_rejection(rejection.mappingStatus);
        }
        return {.reason = mean_field::operators::StellarEquilibriumPreparationRejectionReason::non_finite_physics};
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection
    make_barotropic_rejection(const mean_field::operators::BarotropicClosurePreparationRejection &rejection) {
        using ChildReason = mean_field::operators::BarotropicClosurePreparationRejectionReason;
        switch (rejection.reason) {
        case ChildReason::mapping_failure:
            return make_mapping_rejection(rejection.mappingStatus);
        case ChildReason::equation_of_state:
            return make_thermodynamic_rejection(rejection.equationOfStateError);
        case ChildReason::invalid_quadrature_data:
            return {.reason = mean_field::operators::StellarEquilibriumPreparationRejectionReason::non_finite_physics};
        }
        throw std::logic_error("An unknown barotropic trial rejection reached the stellar root.");
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection
    make_displacement_rejection(const mean_field::operators::DisplacementResidualPreparationRejection &rejection) {
        using ChildReason     = mean_field::operators::DisplacementResidualPreparationRejectionReason;
        using ChildSource     = mean_field::operators::DisplacementResidualPreparationRejectionSource;
        using RootStage       = mean_field::operators::StellarEquilibriumPreparationStage;

        const RootStage stage = [&] {
            switch (rejection.source) {
            case ChildSource::pressure:
                return RootStage::pressure_force;
            case ChildSource::gravity:
                return RootStage::gravity_displacement_force;
            case ChildSource::rotation:
                return RootStage::rotational_displacement_force;
            case ChildSource::composition:
                return RootStage::displacement_composition;
            }
            return RootStage::displacement_residual;
        }();

        switch (rejection.reason) {
        case ChildReason::equation_of_state: {
            auto rootRejection  = make_thermodynamic_rejection(rejection.equationOfStateCode);
            rootRejection.stage = stage;
            return rootRejection;
        }
        case ChildReason::invalid_mapping: {
            auto rootRejection  = make_mapping_rejection(rejection.mappingStatus);
            rootRejection.stage = stage;
            return rootRejection;
        }
        case ChildReason::non_finite_arithmetic:
            return {
                .reason = mean_field::operators::StellarEquilibriumPreparationRejectionReason::non_finite_physics,
                .stage  = stage
            };
        }
        throw std::logic_error("An unknown displacement trial rejection reached the stellar root.");
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection
    make_mass_rejection(const mean_field::operators::MassNormalizationPreparationRejection &rejection) {
        using ChildReason = mean_field::operators::MassNormalizationPreparationRejectionReason;
        if (rejection.reason == ChildReason::mapping_failure) {
            return make_mapping_rejection(rejection.mappingStatus);
        }
        return {.reason = mean_field::operators::StellarEquilibriumPreparationRejectionReason::non_finite_physics};
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumPreparationRejection
    make_hydrostatic_rejection(const mean_field::operators::HydrostaticEquilibriumPreparationRejection &rejection) {
        using ChildReason = mean_field::operators::HydrostaticEquilibriumPreparationRejectionReason;
        using Failure     = mean_field::operators::StellarEquilibriumPreparationRejection;
        using Reason      = mean_field::operators::StellarEquilibriumPreparationRejectionReason;

        switch (rejection.reason) {
        case ChildReason::inverted_geometry:
            return Failure{.reason = Reason::inverted_geometry};
        case ChildReason::non_finite_geometry:
            return Failure{.reason = Reason::non_finite_geometry};
        case ChildReason::non_finite_residual:
            return Failure{.reason = Reason::non_finite_physics};
        }
        throw std::logic_error("An unknown hydrostatic trial rejection reached the stellar root.");
    }

    [[nodiscard]] std::array<
        int,
        StellarRootForm::value_block_count>
    make_value_sizes(
        const mean_field::field::FieldDofMap &densityMap,
        const int surfaceDeformationParameterCount,
        const mean_field::field::FieldDofMap &gravityFluxMap,
        const mean_field::field::FieldDofMap &gravityPotentialMap,
        const mean_field::field::FieldDofMap &enthalpyMap
    ) {
        return {densityMap.reduced_size(),          surfaceDeformationParameterCount, gravityFluxMap.reduced_size(),
                gravityPotentialMap.reduced_size(), enthalpyMap.reduced_size(),       1};
    }

    [[nodiscard]] std::array<
        int,
        StellarRootForm::residual_block_count>
    make_residual_sizes(
        const mean_field::field::FieldDofMap &densityMap,
        const int surfaceDeformationParameterCount,
        const mean_field::field::FieldDofMap &gravityFluxMap,
        const mean_field::field::FieldDofMap &gravityPotentialMap,
        const mean_field::field::FieldDofMap &enthalpyMap
    ) {
        return {gravityFluxMap.reduced_size(),    gravityPotentialMap.reduced_size(), densityMap.reduced_size(),
                surfaceDeformationParameterCount, enthalpyMap.reduced_size(),         1};
    }

    [[nodiscard]] mfem::Array<int> make_gravity_state_offsets(
        const mean_field::field::FieldDofMap &densityMap,
        const mean_field::field::FieldDofMap &displacementMap,
        const mean_field::field::FieldDofMap &gravityFluxMap,
        const mean_field::field::FieldDofMap &gravityPotentialMap
    ) {
        mfem::Array<int> offsets(5);
        offsets[0] = 0;
        offsets[1] = offsets[0] + densityMap.reduced_size();
        offsets[2] = offsets[1] + displacementMap.reduced_size();
        offsets[3] = offsets[2] + gravityFluxMap.reduced_size();
        offsets[4] = offsets[3] + gravityPotentialMap.reduced_size();
        return offsets;
    }

    [[nodiscard]] mfem::Array<int> make_gravity_residual_offsets(
        const mean_field::field::FieldDofMap &gravityFluxMap,
        const mean_field::field::FieldDofMap &gravityPotentialMap
    ) {
        mfem::Array<int> offsets(3);
        offsets[0] = 0;
        offsets[1] = gravityFluxMap.reduced_size();
        offsets[2] = offsets[1] + gravityPotentialMap.reduced_size();
        return offsets;
    }

    void assign_gravity_block(
        mfem::Vector &gravityState,
        const mfem::Array<int> &offsets,
        const int blockIndex,
        const mfem::Vector &source,
        const char *message
    ) {
        MFEM_VERIFY(offsets.Size() == 5, "Gravity state offsets are invalid.");
        MFEM_VERIFY(blockIndex >= 0 && blockIndex + 1 < offsets.Size(), "Requested gravity-state block is invalid.");

        const int blockSize = offsets[blockIndex + 1] - offsets[blockIndex];
        MFEM_VERIFY(blockSize == source.Size(), message);
        MFEM_VERIFY(gravityState.Size() == offsets.Last(), "Packed gravity state has the wrong size.");

        mfem::Vector destination(gravityState.GetData() + offsets[blockIndex], blockSize);
        destination = source;
    }

    void pack_gravity_vector(
        mfem::Vector &gravityState,
        const mfem::Array<int> &offsets,
        const mfem::Vector &density,
        const mfem::Vector &displacement,
        const mfem::Vector &gravityGradient,
        const mfem::Vector &gravityPotential
    ) {
        MFEM_VERIFY(offsets.Size() == 5, "Packed gravity state requires four blocks.");
        if (gravityState.Size() != offsets.Last()) {
            gravityState.SetSize(offsets.Last());
        }

        assign_gravity_block(
            gravityState, offsets, 0, density, "The full density vector has the wrong gravity-state size."
        );
        assign_gravity_block(
            gravityState, offsets, 1, displacement, "The displacement vector has the wrong gravity-state size."
        );
        assign_gravity_block(
            gravityState, offsets, 2, gravityGradient, "The gravity-gradient vector has the wrong gravity-state size."
        );
        assign_gravity_block(
            gravityState, offsets, 3, gravityPotential, "The gravity-potential vector has the wrong gravity-state size."
        );
    }

    void validate_finite_vector(
        const mfem::Vector &vector,
        const char *message
    ) {
        for (int index = 0; index < vector.Size(); ++index) {
            MFEM_VERIFY(std::isfinite(vector(index)), message);
        }
    }

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

    void validate_dependency_transition(
        const mean_field::operators::StellarEquilibriumDependencyStamp &prepared,
        const mean_field::operators::StellarEquilibriumDependencyStamp &requested,
        const char *message
    ) {
        MFEM_VERIFY(prepared.identity != requested.identity || requested.revision >= prepared.revision, message);
        MFEM_VERIFY(
            prepared.identity == requested.identity || prepared.revision != requested.revision,
            "A new stellar-equilibrium dependency identity must also carry a visibly different revision."
        );
    }

    [[nodiscard]] mean_field::operators::context::gravity_field::GravityFieldRevisions make_gravity_revisions(
        const mean_field::operators::StellarEquilibriumDependencies &dependencies,
        const mean_field::operators::StellarEquilibriumDependencyStamp &generatedDisplacement
    ) {
        return {
            .discretization    = {.value = dependencies.discretization.revision},
            .displacement      = {.value = generatedDisplacement.revision},
            .density           = {.value = dependencies.density.revision},
            .gravity_gradient  = {.value = dependencies.gravityGradient.revision},
            .gravity_potential = {.value = dependencies.gravityPotential.revision}
        };
    }

    [[nodiscard]] mean_field::operators::context::barotropic::BarotropicClosureDependencies
    make_barotropic_closure_dependencies(
        const mean_field::operators::StellarEquilibriumDependencies &dependencies,
        const mean_field::operators::StellarEquilibriumDependencyStamp &generatedDisplacement
    ) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .density      = {.identity = dependencies.density.identity, .revision = dependencies.density.revision},
            .enthalpy     = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision},
            .displacement = {.identity = generatedDisplacement.identity, .revision = generatedDisplacement.revision}
        };
    }

    [[nodiscard]] mean_field::operators::DisplacementResidualDependencies make_displacement_dependencies(
        const mean_field::operators::StellarEquilibriumDependencies &dependencies,
        const mean_field::operators::StellarEquilibriumDependencyStamp &generatedDisplacement
    ) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .density      = {.identity = dependencies.density.identity, .revision = dependencies.density.revision},
            .displacement = {.identity = generatedDisplacement.identity, .revision = generatedDisplacement.revision},
            .gravityGradient =
                {.identity = dependencies.gravityGradient.identity, .revision = dependencies.gravityGradient.revision},
            .enthalpy = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision},
            .rotation = {.identity = dependencies.rotation.identity, .revision = dependencies.rotation.revision}
        };
    }

    [[nodiscard]] mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies
    make_hydrostatic_dependencies(
        const mean_field::operators::StellarEquilibriumDependencies &dependencies,
        const mean_field::operators::StellarEquilibriumDependencyStamp &generatedDisplacement
    ) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .enthalpy = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision},
            .gravityPotential =
                {.identity = dependencies.gravityPotential.identity,
                 .revision = dependencies.gravityPotential.revision},
            .displacement = {.identity = generatedDisplacement.identity, .revision = generatedDisplacement.revision},
            .rotation     = {.identity = dependencies.rotation.identity, .revision = dependencies.rotation.revision},
            .bernoulliConstant = {
                .identity = dependencies.bernoulliConstant.identity, .revision = dependencies.bernoulliConstant.revision
            }
        };
    }

    [[nodiscard]] mean_field::operators::MassNormalizationDependencies make_mass_dependencies(
        const mean_field::operators::StellarEquilibriumDependencies &dependencies,
        const mean_field::operators::StellarEquilibriumDependencyStamp &generatedDisplacement
    ) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .density      = {.identity = dependencies.density.identity, .revision = dependencies.density.revision},
            .displacement = {.identity = generatedDisplacement.identity, .revision = generatedDisplacement.revision},
            .targetMass   = {.identity = dependencies.targetMass.identity, .revision = dependencies.targetMass.revision}
        };
    }
} // namespace

namespace mean_field::operators {
    struct PreparedStellarEquilibriumOperator::ConstructionData {
        deformation::PreparedDomainDeformationRuntime domainDeformation;
        field::FieldDofMap densityMap;
        field::FieldDofMap displacementMap;
        field::FieldDofMap gravityFluxMap;
        field::FieldDofMap gravityPotentialMap;
        field::FieldDofMap enthalpyMap;
        field::FieldBoundaryDofMap pressureSurfaceRows;

        std::array<int, StellarRootForm::value_block_count> valueSizes;
        std::array<int, StellarRootForm::residual_block_count> residualSizes;
        mfem::Array<int> gravityStateOffsets;
        mfem::Array<int> gravityResidualOffsets;

        ConstructionData(
            fem::FEM &f,
            deformation::PreparedDomainDeformationRuntime preparedDomainDeformation
        )
            : domainDeformation(std::move(preparedDomainDeformation)),
              densityMap(
                  field::make_field_dof_map<
                      field::Density,
                      DomainSchema>(*f.densityFes)
              ),
              displacementMap(
                  field::make_field_dof_map<
                      field::Displacement,
                      DomainSchema>(*f.displacementFes)
              ),
              gravityFluxMap(
                  field::make_field_dof_map<
                      field::Gravity,
                      DomainSchema>(*f.gravityFluxFes)
              ),
              gravityPotentialMap(
                  field::make_field_dof_map<
                      field::Gravity,
                      DomainSchema>(*f.gravityPotentialFes)
              ),
              enthalpyMap(
                  field::make_field_dof_map<
                      field::Enthalpy,
                      DomainSchema>(*f.enthalpyFes)
              ),
              pressureSurfaceRows(
                  field::make_field_boundary_dof_map<
                      field::Enthalpy,
                      utils::domain::StellarSurface,
                      DomainSchema>(
                      *f.enthalpyFes,
                      enthalpyMap
                  )
              ),
              valueSizes(make_value_sizes(
                  densityMap,
                  domainDeformation.parameterCount(),
                  gravityFluxMap,
                  gravityPotentialMap,
                  enthalpyMap
              )),
              residualSizes(make_residual_sizes(
                  densityMap,
                  domainDeformation.parameterCount(),
                  gravityFluxMap,
                  gravityPotentialMap,
                  enthalpyMap
              )),
              gravityStateOffsets(make_gravity_state_offsets(
                  densityMap,
                  displacementMap,
                  gravityFluxMap,
                  gravityPotentialMap
              )),
              gravityResidualOffsets(make_gravity_residual_offsets(
                  gravityFluxMap,
                  gravityPotentialMap
              )) {
        }
    };

    PreparedStellarEquilibriumOperator::ConstructionData PreparedStellarEquilibriumOperator::MakeConstructionData(
        fem::FEM &f,
        deformation::PreparedDomainDeformationRuntime domainDeformation
    ) {
        verify_coupled_discretization(f);
        return ConstructionData(f, std::move(domainDeformation));
    }

    PreparedStellarEquilibriumOperator::PreparedStellarEquilibriumOperator(
        fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &equationOfState,
        models::CompiledFixedMass fixedMassConstraint,
        const PressureSurfaceConstraintView surfaceConstraint,
        deformation::PreparedDomainDeformationRuntime domainDeformation
    )
        : PreparedStellarEquilibriumOperator(
              f,
              domainMapper,
              equationOfState,
              std::move(fixedMassConstraint),
              surfaceConstraint,
              MakeConstructionData(
                  f,
                  std::move(domainDeformation)
              )
          ) {
    }

    PreparedStellarEquilibriumOperator::PreparedStellarEquilibriumOperator(
        fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &equationOfState,
        models::CompiledFixedMass fixedMassConstraint,
        const PressureSurfaceConstraintView surfaceConstraint,
        ConstructionData constructionData
    )
        : mfem::Operator(
              StellarEquilibriumLayout(
                  constructionData.valueSizes,
                  constructionData.residualSizes
              )
                  .residual_offsets()
                  .Last(),
              StellarEquilibriumLayout(
                  constructionData.valueSizes,
                  constructionData.residualSizes
              )
                  .value_offsets()
                  .Last()
          ),
          m_rootManifest(
              constructionData.valueSizes,
              constructionData.residualSizes,
              StellarEquilibriumSpecificationModel{
                  equationOfState,
                  surface::Isobaric{dimensions::PressureValue{surfaceConstraint.descriptor().targetPressure}},
                  fixedMassConstraint.specification()
              },
              constructionData.pressureSurfaceRows.size()
          ),
          m_communicator(f.mesh->GetComm()),
          m_gravityStateOffsets(constructionData.gravityStateOffsets),
          m_gravityContext(
              f,
              domainMapper
          ),
          m_gravityJacobianOperator(
              f,
              domainMapper,
              m_gravityContext,
              m_gravityStateOffsets,
              constructionData.gravityResidualOffsets
          ),
          m_gravityOperator(
              f,
              domainMapper,
              m_gravityContext,
              m_gravityStateOffsets,
              m_gravityJacobianOperator
          ),
          m_barotropicClosureOperator(
              f,
              domainMapper,
              equationOfState
          ),
          m_hydrostaticOperator(
              f,
              domainMapper
          ),
          m_displacementOperator(
              f,
              domainMapper,
              equationOfState,
              m_gravityContext
          ),
          m_massNormalizationOperator(
              f,
              domainMapper,
              m_gravityContext
          ),
          m_surfaceConstraintOperator(
              constructionData.pressureSurfaceRows,
              surfaceConstraint
          ),
          m_domainDeformation(std::move(constructionData.domainDeformation)),
          m_fixedMassConstraint(std::move(fixedMassConstraint)) {

        MFEM_VERIFY(
            Width() == m_rootManifest.layout().value_offsets().Last() &&
                Height() == m_rootManifest.layout().residual_offsets().Last(),
            "PreparedStellarEquilibriumOperator has inconsistent block dimensions."
        );

        MFEM_VERIFY(
            m_domainDeformation.volumeDisplacementSize() == constructionData.displacementMap.reduced_size(),
            "The domain-deformation output does not match the coupled displacement discretization."
        );

        m_generatedDisplacementDependency.identity =
            static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(&m_domainDeformation));

        m_gravityState.SetSize(m_gravityStateOffsets.Last());

        m_gravityDirection.SetSize(m_gravityStateOffsets.Last());

        m_surfaceDeformationParameters.SetSize(m_domainDeformation.parameterCount());
        m_generatedVolumeDisplacement.SetSize(m_domainDeformation.volumeDisplacementSize());
        m_fullMechanicalResidual.SetSize(m_domainDeformation.volumeDisplacementSize());
        m_volumeDisplacementDirection.SetSize(m_domainDeformation.volumeDisplacementSize());
        m_fullMechanicalAction.SetSize(m_domainDeformation.volumeDisplacementSize());
        m_surfaceShapeAction.SetSize(m_domainDeformation.parameterCount());
        m_pullbackDerivativeAction.SetSize(m_domainDeformation.parameterCount());
        m_densityVolumeIntegralAction.SetSize(1);

        m_gravityState                 = 0.0;
        m_gravityDirection             = 0.0;
        m_surfaceDeformationParameters = 0.0;
        m_generatedVolumeDisplacement  = 0.0;
        m_fullMechanicalResidual       = 0.0;
        m_volumeDisplacementDirection  = 0.0;
        m_fullMechanicalAction         = 0.0;
        m_surfaceShapeAction           = 0.0;
        m_pullbackDerivativeAction     = 0.0;
        m_densityVolumeIntegralAction  = 0.0;
    }

    PreparedStellarEquilibriumReport PreparedStellarEquilibriumOperator::Prepare(
        const mfem::Vector &state,
        const StellarEquilibriumDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        auto result = TryPrepare(state, dependencies, rotation);
        if (!result.has_value()) {
            throwStellarEquilibriumPreparationRejection(result.error());
        }
        return std::move(result).value();
    }

    StellarEquilibriumPreparationResult<PreparedStellarEquilibriumReport>
    PreparedStellarEquilibriumOperator::TryPrepare(
        const mfem::Vector &state,
        const StellarEquilibriumDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        MFEM_VERIFY(
            state.Size() == Width(), "PreparedStellarEquilibriumOperator received a state with the wrong size."
        );

        const int localStateIsFinite = vector_is_finite(state) ? 1 : 0;
        int globalStateIsFinite      = 0;
        if (MPI_Allreduce(&localStateIsFinite, &globalStateIsFinite, 1, MPI_INT, MPI_MIN, m_communicator) !=
            MPI_SUCCESS) {
            throw std::runtime_error("PreparedStellarEquilibriumOperator could not synchronize state validity.");
        }
        if (globalStateIsFinite == 0) {
            m_isPrepared = false;
            return std::unexpected(
                StellarEquilibriumPreparationRejection{
                    .reason = StellarEquilibriumPreparationRejectionReason::non_finite_physics
                }
            );
        }

        const bool wasPrepared = m_isPrepared;
        if (wasPrepared) {
            validate_dependency_transition(
                m_preparedDependencies.discretization, dependencies.discretization,
                "The discretization revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.density, dependencies.density, "The density revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.surfaceDeformation, dependencies.surfaceDeformation,
                "The surface-deformation revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.gravityGradient, dependencies.gravityGradient,
                "The gravity-gradient revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.gravityPotential, dependencies.gravityPotential,
                "The gravity-potential revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.enthalpy, dependencies.enthalpy, "The enthalpy revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.bernoulliConstant, dependencies.bernoulliConstant,
                "The Bernoulli-constant revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.rotation, dependencies.rotation, "The rotation revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.targetMass, dependencies.targetMass,
                "The target-mass revision cannot move backwards."
            );
        }

        m_isPrepared              = false;

        const auto rootState      = m_rootManifest.stateView(state);

        const auto reducedDensity = rootState.block(utils::blocks::density_field.mass_term);
        const auto surfaceDeformationParameters =
            rootState.block(utils::blocks::surface_deformation_field.parameters_term);
        const auto gravityGradient  = rootState.block(utils::blocks::gravity_field.gradient_term);
        const auto gravityPotential = rootState.block(utils::blocks::gravity_field.poisson_term);
        const auto reducedEnthalpy  = rootState.block(utils::blocks::enthalpy_field.specific_term);
        const auto bernoulli = rootState.block(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);

        const bool generatedGeometryChanged =
            !wasPrepared || dependencies.discretization != m_preparedDependencies.discretization ||
            dependencies.surfaceDeformation != m_preparedDependencies.surfaceDeformation;

        PreparedStellarEquilibriumReport report;
        if (generatedGeometryChanged) {
            m_surfaceDeformationParameters = surfaceDeformationParameters;
            m_domainDeformation.buildVolumeDisplacement(m_surfaceDeformationParameters, m_generatedVolumeDisplacement);
            const deformation::DomainDeformationGeometryReport generatedGeometry =
                m_domainDeformation.inspectMappedGeometry(m_generatedVolumeDisplacement);
            if (!std::isfinite(generatedGeometry.minimumJacobianDeterminant)) {
                return std::unexpected(
                    StellarEquilibriumPreparationRejection{
                        .reason                     = StellarEquilibriumPreparationRejectionReason::non_finite_geometry,
                        .stage                      = StellarEquilibriumPreparationStage::generated_geometry,
                        .minimumJacobianDeterminant = generatedGeometry.minimumJacobianDeterminant
                    }
                );
            }
            if (!generatedGeometry.isOrientationPreserving()) {
                return std::unexpected(
                    StellarEquilibriumPreparationRejection{
                        .reason                     = StellarEquilibriumPreparationRejectionReason::inverted_geometry,
                        .stage                      = StellarEquilibriumPreparationStage::generated_geometry,
                        .minimumJacobianDeterminant = generatedGeometry.minimumJacobianDeterminant
                    }
                );
            }
            m_generatedGeometryReport = generatedGeometry;
            ++m_generatedDisplacementDependency.revision;
            ++m_statistics.generatedGeometryBuilds;
            report.generatedVolumeDisplacement = true;
        }
        report.generatedGeometry     = m_generatedGeometryReport;
        report.generatedDisplacement = m_generatedDisplacementDependency;

        pack_gravity_vector(
            m_gravityState, m_gravityStateOffsets, reducedDensity, m_generatedVolumeDisplacement, gravityGradient,
            gravityPotential
        );

        auto gravityResult = m_gravityOperator.TryPrepare(
            m_gravityState, make_gravity_revisions(dependencies, m_generatedDisplacementDependency)
        );
        if (!gravityResult.has_value()) {
            return std::unexpected(with_preparation_stage(
                make_gravity_rejection(gravityResult.error()), StellarEquilibriumPreparationStage::gravity
            ));
        }
        report.gravity          = std::move(gravityResult).value();

        /*
         * Mechanical-force preparation consumes the shared gravity context,
         * but it is independent of the closure and hydrostatic rows.  Prepare
         * it as soon as that dependency is ready so a mapped-force rejection
         * does not pay for unrelated candidate rows first.
         */
        auto displacementResult = m_displacementOperator.TryPrepare(
            {.enthalpy = reducedEnthalpy},
            make_displacement_dependencies(dependencies, m_generatedDisplacementDependency), rotation
        );
        if (!displacementResult.has_value()) {
            return std::unexpected(make_displacement_rejection(displacementResult.error()));
        }
        report.displacement          = std::move(displacementResult).value();

        auto barotropicClosureResult = m_barotropicClosureOperator.TryPrepare(
            {.density = reducedDensity, .enthalpy = reducedEnthalpy, .displacement = m_generatedVolumeDisplacement},
            make_barotropic_closure_dependencies(dependencies, m_generatedDisplacementDependency)
        );
        if (!barotropicClosureResult.has_value()) {
            return std::unexpected(with_preparation_stage(
                make_barotropic_rejection(barotropicClosureResult.error()),
                StellarEquilibriumPreparationStage::barotropic_closure
            ));
        }
        report.barotropicClosure = std::move(barotropicClosureResult).value();

        auto hydrostaticResult   = m_hydrostaticOperator.TryPrepare(
            {.enthalpy          = reducedEnthalpy,
             .gravityPotential  = gravityPotential,
             .displacement      = m_generatedVolumeDisplacement,
             .bernoulliConstant = bernoulli(0)},
            make_hydrostatic_dependencies(dependencies, m_generatedDisplacementDependency), rotation
        );
        if (!hydrostaticResult.has_value()) {
            return std::unexpected(with_preparation_stage(
                make_hydrostatic_rejection(hydrostaticResult.error()),
                StellarEquilibriumPreparationStage::hydrostatic_equilibrium
            ));
        }
        report.hydrostatic           = std::move(hydrostaticResult).value();

        auto massNormalizationResult = m_massNormalizationOperator.TryPrepare(
            m_fixedMassConstraint, make_mass_dependencies(dependencies, m_generatedDisplacementDependency)
        );
        if (!massNormalizationResult.has_value()) {
            return std::unexpected(with_preparation_stage(
                make_mass_rejection(massNormalizationResult.error()),
                StellarEquilibriumPreparationStage::mass_normalization
            ));
        }
        report.massNormalization = std::move(massNormalizationResult).value();

        report.surfaceConstraint = m_surfaceConstraintOperator.Prepare(
            reducedEnthalpy, !wasPrepared || dependencies.enthalpy != m_preparedDependencies.enthalpy
        );

        const bool dependenciesChanged = !wasPrepared || dependencies != m_preparedDependencies;
        if (dependenciesChanged || report.DidAnyChildWork()) {
            AssembleResidual();
            report.assembledResidual = true;
        }

        m_preparedDependencies = dependencies;
        m_isPrepared           = true;
        return report;
    }

    void PreparedStellarEquilibriumOperator::AssembleResidual() {
        mfem::Vector gravity;
        mfem::Vector closure;
        mfem::Vector surfaceShape;
        mfem::Vector hydrostatic;
        mfem::Vector mass;

        m_gravityOperator.Mult(m_gravityState, gravity);
        m_barotropicClosureOperator.BuildResidual(closure);
        m_displacementOperator.BuildResidual(m_fullMechanicalResidual);
        surfaceShape.SetSize(m_domainDeformation.parameterCount());
        m_domainDeformation.applyJacobianTranspose(
            m_surfaceDeformationParameters, m_fullMechanicalResidual, surfaceShape
        );
        m_hydrostaticOperator.BuildResidual(hydrostatic);
        m_surfaceConstraintOperator.ApplyResidualRows(hydrostatic);
        m_massNormalizationOperator.BuildResidual(mass);

        m_cachedResidual.SetSize(Height());
        m_cachedResidual        = 0.0;
        const auto residualView = m_rootManifest.residualView(m_cachedResidual);

        MFEM_VERIFY(
            gravity.Size() == residualView.block(utils::blocks::gravity_field.gradient_term).Size() +
                                  residualView.block(utils::blocks::gravity_field.poisson_term).Size(),
            "The gravity residual has the wrong size."
        );

        mfem::Vector gravityGradient(
            gravity.GetData(), residualView.block(utils::blocks::gravity_field.gradient_term).Size()
        );
        mfem::Vector gravityPotential(
            gravity.GetData() + gravityGradient.Size(),
            residualView.block(utils::blocks::gravity_field.poisson_term).Size()
        );

        residualView.assign(utils::blocks::gravity_field.gradient_term, gravityGradient);
        residualView.assign(utils::blocks::gravity_field.poisson_term, gravityPotential);
        residualView.assign(utils::blocks::density_field.mass_term, closure);
        residualView.assign(utils::blocks::surface_deformation_field.shape_equilibrium_term, surfaceShape);
        residualView.assign(utils::blocks::enthalpy_field.specific_term, hydrostatic);
        residualView.assign(utils::blocks::fixed_total_mass_constraint.mass_normalization_term, mass);

        ++m_statistics.residualAssemblies;
    }

    void PreparedStellarEquilibriumOperator::BuildResidual(mfem::Vector &residual) const {
        VerifyPrepared();
        residual = m_cachedResidual;
        ++m_statistics.residualApplications;
    }

    void PreparedStellarEquilibriumOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(
            direction.Size() == Width(),
            "PreparedStellarEquilibriumOperator received a Jacobian direction with the wrong size."
        );
        validate_finite_vector(
            direction, "PreparedStellarEquilibriumOperator received a non-finite Jacobian direction."
        );

        const auto rootDirection           = m_rootManifest.directionView(direction);

        const auto reducedDensityDirection = rootDirection.block(utils::blocks::density_field.mass_term);
        const auto surfaceDeformationDirection =
            rootDirection.block(utils::blocks::surface_deformation_field.parameters_term);
        const auto gravityGradientDirection  = rootDirection.block(utils::blocks::gravity_field.gradient_term);
        const auto gravityPotentialDirection = rootDirection.block(utils::blocks::gravity_field.poisson_term);
        const auto reducedEnthalpyDirection  = rootDirection.block(utils::blocks::enthalpy_field.specific_term);
        const auto bernoulliDirection =
            rootDirection.block(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);

        m_domainDeformation.applyJacobian(
            m_surfaceDeformationParameters, surfaceDeformationDirection, m_volumeDisplacementDirection
        );

        pack_gravity_vector(
            m_gravityDirection, m_gravityStateOffsets, reducedDensityDirection, m_volumeDisplacementDirection,
            gravityGradientDirection, gravityPotentialDirection
        );

        mfem::Vector gravityAction;
        mfem::Vector closureAction;
        mfem::Vector hydrostaticAction;
        mfem::Vector massAction;

        m_gravityJacobianOperator.Mult(m_gravityDirection, gravityAction);

        m_barotropicClosureOperator.Mult(
            reducedDensityDirection, reducedEnthalpyDirection, m_volumeDisplacementDirection, closureAction
        );

        m_displacementOperator.ApplyCompleteJacobianAction(
            reducedDensityDirection, m_volumeDisplacementDirection, gravityGradientDirection, reducedEnthalpyDirection,
            m_fullMechanicalAction
        );
        m_domainDeformation.applyJacobianTranspose(
            m_surfaceDeformationParameters, m_fullMechanicalAction, m_surfaceShapeAction
        );
        m_domainDeformation.applyPullbackDerivative(
            m_surfaceDeformationParameters, surfaceDeformationDirection, m_fullMechanicalResidual,
            m_pullbackDerivativeAction
        );
        m_surfaceShapeAction += m_pullbackDerivativeAction;

        m_hydrostaticOperator.ApplyCompleteJacobianAction(
            reducedEnthalpyDirection, gravityPotentialDirection, bernoulliDirection(0), m_volumeDisplacementDirection,
            hydrostaticAction
        );
        m_surfaceConstraintOperator.ApplyJacobianRows(reducedEnthalpyDirection, hydrostaticAction);

        m_massNormalizationOperator.ApplyCompleteJacobianAction(
            reducedDensityDirection, m_volumeDisplacementDirection, massAction
        );

        action.SetSize(Height());
        action                = 0.0;
        const auto actionView = m_rootManifest.residualView(action);

        MFEM_VERIFY(
            gravityAction.Size() == actionView.block(utils::blocks::gravity_field.gradient_term).Size() +
                                        actionView.block(utils::blocks::gravity_field.poisson_term).Size(),
            "The gravity Jacobian action has the wrong size."
        );

        mfem::Vector gravityGradientAction(
            gravityAction.GetData(), actionView.block(utils::blocks::gravity_field.gradient_term).Size()
        );
        mfem::Vector gravityPotentialAction(
            gravityAction.GetData() + gravityGradientAction.Size(),
            actionView.block(utils::blocks::gravity_field.poisson_term).Size()
        );

        actionView.assign(utils::blocks::gravity_field.gradient_term, gravityGradientAction);
        actionView.assign(utils::blocks::gravity_field.poisson_term, gravityPotentialAction);
        actionView.assign(utils::blocks::density_field.mass_term, closureAction);
        actionView.assign(utils::blocks::surface_deformation_field.shape_equilibrium_term, m_surfaceShapeAction);
        actionView.assign(utils::blocks::enthalpy_field.specific_term, hydrostaticAction);
        actionView.assign(utils::blocks::fixed_total_mass_constraint.mass_normalization_term, massAction);

        ++m_statistics.jacobianApplications;
    }

    bool PreparedStellarEquilibriumOperator::IsPrepared() const noexcept {
        return m_isPrepared && m_gravityContext.IsPrepared() && m_barotropicClosureOperator.IsPrepared() &&
               m_hydrostaticOperator.IsPrepared() && m_displacementOperator.IsPrepared() &&
               m_massNormalizationOperator.IsPrepared() && m_surfaceConstraintOperator.IsPrepared();
    }

    double PreparedStellarEquilibriumOperator::GetTargetMass() const noexcept {
        return m_fixedMassConstraint.targetMass().value();
    }

    const StellarEquilibriumLayout &PreparedStellarEquilibriumOperator::GetLayout() const noexcept {
        return m_rootManifest.layout();
    }

    const StellarEquilibriumRootManifest &PreparedStellarEquilibriumOperator::GetRootManifest() const noexcept {
        return m_rootManifest;
    }

    RootStateView<utils::blocks::surface_deformed_stellar_equilibrium_form>
    PreparedStellarEquilibriumOperator::GetRootStateView(const mfem::Vector &state) const {
        return m_rootManifest.stateView(state);
    }

    ResidualView<utils::blocks::surface_deformed_stellar_equilibrium_form>
    PreparedStellarEquilibriumOperator::GetResidualView(mfem::Vector &residual) const {
        return m_rootManifest.residualView(residual);
    }

    RootConstraintReport PreparedStellarEquilibriumOperator::GetFixedMassReport() const {
        VerifyPrepared();
        return m_rootManifest.fixedMassReport(m_massNormalizationOperator.GetCurrentMass());
    }

    const StellarEquilibriumDependencies &PreparedStellarEquilibriumOperator::GetDependencies() const {
        VerifyPrepared();
        return m_preparedDependencies;
    }

    const PreparedStellarEquilibriumStatistics &PreparedStellarEquilibriumOperator::GetStatistics() const noexcept {
        return m_statistics;
    }

    const context::gravity_field::GravityFieldLinearizationContext &
    PreparedStellarEquilibriumOperator::GetGravityContext() const noexcept {
        return m_gravityContext;
    }

    const GravityFieldOperator &PreparedStellarEquilibriumOperator::GetGravityOperator() const noexcept {
        return m_gravityOperator;
    }

    const GravityFieldJacobianOperator &
    PreparedStellarEquilibriumOperator::GetGravityJacobianOperator() const noexcept {
        return m_gravityJacobianOperator;
    }

    const PreparedBarotropicClosureOperator &
    PreparedStellarEquilibriumOperator::GetBarotropicClosureOperator() const noexcept {
        return m_barotropicClosureOperator;
    }

    const context::barotropic::BarotropicClosureLinearizationContext &
    PreparedStellarEquilibriumOperator::GetBarotropicClosureContext() const noexcept {
        return m_barotropicClosureOperator.GetContext();
    }

    const PreparedHydrostaticEquilibriumOperator &
    PreparedStellarEquilibriumOperator::GetHydrostaticOperator() const noexcept {
        return m_hydrostaticOperator;
    }

    const PreparedDisplacementResidualOperator &
    PreparedStellarEquilibriumOperator::GetDisplacementOperator() const noexcept {
        return m_displacementOperator;
    }

    const PreparedMassNormalizationOperator &
    PreparedStellarEquilibriumOperator::GetMassNormalizationOperator() const noexcept {
        return m_massNormalizationOperator;
    }

    double PreparedStellarEquilibriumOperator::ApplyDensityVolumeIntegralDensityAction(
        const mfem::Vector &densityDirection
    ) const {
        VerifyPrepared();
        m_massNormalizationOperator.ApplyDensityJacobianAction(densityDirection, m_densityVolumeIntegralAction);
        MFEM_VERIFY(
            m_densityVolumeIntegralAction.Size() == 1, "The density-volume integral must produce one global scalar."
        );
        return m_densityVolumeIntegralAction(0);
    }

    double PreparedStellarEquilibriumOperator::ApplyDensityVolumeIntegralSurfaceShapeAction(
        const mfem::Vector &surfaceShapeDirection
    ) const {
        VerifyPrepared();
        m_domainDeformation.applyJacobian(
            m_surfaceDeformationParameters, surfaceShapeDirection, m_volumeDisplacementDirection
        );
        m_massNormalizationOperator.ApplyDisplacementJacobianAction(
            m_volumeDisplacementDirection, m_densityVolumeIntegralAction
        );
        MFEM_VERIFY(
            m_densityVolumeIntegralAction.Size() == 1,
            "The density-volume shape derivative must produce one global scalar."
        );
        return m_densityVolumeIntegralAction(0);
    }

    const PreparedPressureSurfaceConstraint &
    PreparedStellarEquilibriumOperator::GetSurfaceConstraintOperator() const noexcept {
        return m_surfaceConstraintOperator;
    }

    const deformation::PreparedDomainDeformationRuntime &
    PreparedStellarEquilibriumOperator::GetDomainDeformation() const noexcept {
        return m_domainDeformation;
    }

    const mfem::Vector &PreparedStellarEquilibriumOperator::GetSurfaceDeformationParameters() const {
        VerifyPrepared();
        return m_surfaceDeformationParameters;
    }

    const mfem::Vector &PreparedStellarEquilibriumOperator::GetGeneratedVolumeDisplacement() const {
        VerifyPrepared();
        return m_generatedVolumeDisplacement;
    }

    void PreparedStellarEquilibriumOperator::BuildVolumeDisplacementDirection(
        const mfem::Vector &surfaceDeformationDirection,
        mfem::Vector &volumeDisplacementDirection
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(
            surfaceDeformationDirection.Size() == m_domainDeformation.parameterCount(),
            "PreparedStellarEquilibriumOperator received a surface-deformation direction with the wrong size."
        );
        validate_finite_vector(
            surfaceDeformationDirection,
            "PreparedStellarEquilibriumOperator received a non-finite surface-deformation direction."
        );
        MFEM_VERIFY(
            volumeDisplacementDirection.Size() == m_domainDeformation.volumeDisplacementSize(),
            "PreparedStellarEquilibriumOperator received a volume-displacement workspace with the wrong size."
        );

        m_domainDeformation.applyJacobian(
            m_surfaceDeformationParameters, surfaceDeformationDirection, volumeDisplacementDirection
        );
    }

    const mfem::Vector &PreparedStellarEquilibriumOperator::GetFullMechanicalResidual() const {
        VerifyPrepared();
        return m_fullMechanicalResidual;
    }

    const StellarEquilibriumDependencyStamp &
    PreparedStellarEquilibriumOperator::GetGeneratedDisplacementDependency() const {
        VerifyPrepared();
        return m_generatedDisplacementDependency;
    }

    void PreparedStellarEquilibriumOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedStellarEquilibriumOperator must be prepared before residual or Jacobian application."
        );
    }
} // namespace mean_field::operators