module;

#include <compare>
#include <concepts>
#include <cstdint>
#include <expected>
#include <limits>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include <mfem.hpp>
#include <mpi.h>

export module mean_field:operators.prepared_stellar_equilibrium;

export import :eos.polytrope;
export import :fem;
export import :field.mfem;
export import :mapping.domain_mapper;
export import :model.stellar;
export import :model.typed_stellar;
export import :operators.context.gravity_field;
export import :operators.gravity_field;
export import :operators.gravity_field_jacobian;
export import :operators.prepared_barotropic_closure;
export import :operators.prepared_displacement_residual;
export import :operators.prepared_hydrostatic_equilibrium;
export import :operators.prepared_mass_normalization;
export import :operators.prepared_surface_constraint;
export import :operators.root_manifest;
export import :physics.rigid_rotation;
export import :utils.blocks;

export namespace mean_field::operators {
    struct StellarEquilibriumDependencyStamp final {
        std::uint64_t identity{0};
        std::uint64_t revision{0};

        constexpr auto operator<=>(const StellarEquilibriumDependencyStamp &) const = default;
    };

    struct StellarEquilibriumDependencies final {
        StellarEquilibriumDependencyStamp discretization;
        StellarEquilibriumDependencyStamp density;
        StellarEquilibriumDependencyStamp surfaceDeformation;
        StellarEquilibriumDependencyStamp gravityGradient;
        StellarEquilibriumDependencyStamp gravityPotential;
        StellarEquilibriumDependencyStamp enthalpy;
        StellarEquilibriumDependencyStamp bernoulliConstant;
        StellarEquilibriumDependencyStamp rotation;
        StellarEquilibriumDependencyStamp targetMass;

        constexpr auto operator<=>(const StellarEquilibriumDependencies &) const = default;
    };

    struct PreparedStellarEquilibriumReport final {
        context::gravity_field::GravityFieldPreparationReport gravity;
        PreparedBarotropicClosureReport barotropicClosure;
        PreparedHydrostaticEquilibriumReport hydrostatic;
        PreparedDisplacementResidualReport displacement;
        PreparedMassNormalizationReport massNormalization;
        PreparedSurfaceConstraintReport surfaceConstraint;
        deformation::DomainDeformationGeometryReport generatedGeometry;
        StellarEquilibriumDependencyStamp generatedDisplacement;
        bool generatedVolumeDisplacement{false};
        bool assembledResidual{false};

        [[nodiscard]] bool DidAnyChildWork() const noexcept {
            return gravity.DidAnyWork() || barotropicClosure.DidAnyWork() || hydrostatic.DidAnyWork() ||
                   displacement.DidAnyWork() || massNormalization.DidAnyWork() || surfaceConstraint.DidAnyWork();
        }

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return DidAnyChildWork() || generatedVolumeDisplacement || assembledResidual;
        }
    };

    enum class StellarEquilibriumPreparationRejectionReason : std::uint8_t {
        inverted_geometry,
        non_finite_geometry,
        thermodynamic_domain,
        non_finite_thermodynamics,
        inadmissible_physics,
        non_finite_physics
    };

    enum class StellarEquilibriumPreparationStage : std::uint8_t {
        unspecified,
        generated_geometry,
        gravity,
        barotropic_closure,
        hydrostatic_equilibrium,
        displacement_residual,
        pressure_force,
        gravity_displacement_force,
        rotational_displacement_force,
        displacement_composition,
        mass_normalization,
        model_specification
    };

    /*
     * A candidate state that cannot define a physical mapped domain is an
     * expected line-search outcome, not an exceptional program failure.  Keep
     * this payload fixed-size so every trial can report it without allocating.
     */
    struct StellarEquilibriumPreparationRejection final {
        StellarEquilibriumPreparationRejectionReason reason{
            StellarEquilibriumPreparationRejectionReason::inverted_geometry
        };
        StellarEquilibriumPreparationStage stage{StellarEquilibriumPreparationStage::unspecified};
        double minimumJacobianDeterminant{std::numeric_limits<double>::quiet_NaN()};
        eos::EvaluationErrorCode thermodynamicErrorCode{eos::EvaluationErrorCode::nonfinite_result};
    };

    template <typename Report>
    using StellarEquilibriumPreparationResult = std::expected<Report, StellarEquilibriumPreparationRejection>;

    [[noreturn]] inline void
    throwStellarEquilibriumPreparationRejection(const StellarEquilibriumPreparationRejection &rejection) {
        switch (rejection.reason) {
        case StellarEquilibriumPreparationRejectionReason::non_finite_geometry:
            throw std::domain_error("The prepared domain deformation has non-finite mapped geometry.");
        case StellarEquilibriumPreparationRejectionReason::thermodynamic_domain:
            throw eos::EvaluationError(
                rejection.thermodynamicErrorCode, "The stellar state lies outside the equation-of-state domain."
            );
        case StellarEquilibriumPreparationRejectionReason::non_finite_thermodynamics:
            throw eos::EvaluationError(
                rejection.thermodynamicErrorCode, "The stellar state produced non-finite thermodynamic data."
            );
        case StellarEquilibriumPreparationRejectionReason::inadmissible_physics:
            throw std::domain_error("The stellar state is physically inadmissible.");
        case StellarEquilibriumPreparationRejectionReason::non_finite_physics:
            throw std::domain_error("The stellar state produced non-finite physical data.");
        case StellarEquilibriumPreparationRejectionReason::inverted_geometry:
            throw std::domain_error("The prepared domain deformation inverts at least one volume element.");
        }
        throw std::logic_error("Unknown stellar-equilibrium candidate-rejection reason.");
    }

    struct PreparedStellarEquilibriumStatistics final {
        std::uint64_t residualAssemblies{0};
        std::uint64_t residualApplications{0};
        std::uint64_t jacobianApplications{0};
        std::uint64_t generatedGeometryBuilds{0};

        constexpr auto operator<=>(const PreparedStellarEquilibriumStatistics &) const = default;
    };

    using StellarEquilibriumLayout =
        utils::blocks::form_layout<utils::blocks::surface_deformed_stellar_equilibrium_form>;

    using StellarEquilibriumSpecificationModel =
        model::StellarModel<models::SpecificationSet<eos::Polytrope, models::FixedTotalMass, surface::Isobaric>>;

    using StellarEquilibriumSystemManifest = EquilibriumSystemManifest<
        StellarEquilibriumSpecificationModel,
        utils::blocks::surface_deformed_stellar_equilibrium_form,
        utils::blocks::surface_deformed_stellar_equilibrium_jacobian_form>;

    using StellarEquilibriumRootManifest = StellarEquilibriumSystemManifest;

    class PreparedStellarEquilibriumOperator final : public mfem::Operator {
    public:
        /*
         * Privileged aggregate runtimes can inspect this complete numerical
         * core. Keep their allow-list on the concrete core itself: a custom
         * EOS may reuse this class, but it cannot extend the class's backend
         * privileges. Ordinary specifications use restricted nested physics
         * and never interact with this list.
         */
        using BackendSpecifications = models::ModelTypeList<
            eos::Polytrope,
            surface::Isobaric,
            models::FixedTotalMass,
            models::FixedAngularMomentum,
            models::FixedCentralDensity>;

        template <models::StellarModelType Model>
            requires std::same_as<
                         typename std::remove_cvref_t<Model>::EquationOfStateType,
                         eos::Polytrope> &&
                     SingleFieldPressureSurfaceConstraintFor<
                         typename std::remove_cvref_t<Model>::SurfaceConstraintType,
                         field::Enthalpy> &&
                     std::is_lvalue_reference_v<Model &&>
        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            Model &&stellarModel
        )
            : PreparedStellarEquilibriumOperator(
                  f,
                  domainMapper,
                  stellarModel.equationOfState(),
                  models::compileConstraint(models::FixedTotalMass{dimensions::MassValue{stellarModel.targetMass()}}),
                  PressureSurfaceConstraintView{stellarModel.compiledSurfaceConstraint()},
                  deformation::PreparedDomainDeformationRuntime{stellarModel.compileDomainDeformation(f)}
              ) {
        }

        /*
         * Authoritative construction path for a compiled equilibrium system.
         * The caller owns the EOS and compiled surface constraint for this
         * operator's lifetime; the remaining compiled contributions are
         * transferred into the operator.
         */
        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            models::CompiledFixedMass fixedMassConstraint,
            PressureSurfaceConstraintView surfaceConstraint,
            deformation::PreparedDomainDeformationRuntime domainDeformation
        );

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

        PreparedStellarEquilibriumReport Prepare(
            const mfem::Vector &state,
            const StellarEquilibriumDependencies &dependencies,
            const physics::RigidRotation &rotation
        );

        [[nodiscard]] StellarEquilibriumPreparationResult<PreparedStellarEquilibriumReport> TryPrepare(
            const mfem::Vector &state,
            const StellarEquilibriumDependencies &dependencies,
            const physics::RigidRotation &rotation
        );

        void BuildResidual(mfem::Vector &residual) const;

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override;

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] double GetTargetMass() const noexcept;
        [[nodiscard]] const StellarEquilibriumLayout &GetLayout() const noexcept;
        [[nodiscard]] const StellarEquilibriumRootManifest &GetRootManifest() const noexcept;
        [[nodiscard]] RootStateView<utils::blocks::surface_deformed_stellar_equilibrium_form>
        GetRootStateView(const mfem::Vector &state) const;
        [[nodiscard]] ResidualView<utils::blocks::surface_deformed_stellar_equilibrium_form>
        GetResidualView(mfem::Vector &residual) const;
        [[nodiscard]] RootConstraintReport GetFixedMassReport() const;
        [[nodiscard]] const StellarEquilibriumDependencies &GetDependencies() const;
        [[nodiscard]] const PreparedStellarEquilibriumStatistics &GetStatistics() const noexcept;

        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetGravityContext() const noexcept;
        [[nodiscard]] const GravityFieldOperator &GetGravityOperator() const noexcept;
        [[nodiscard]] const GravityFieldJacobianOperator &GetGravityJacobianOperator() const noexcept;
        [[nodiscard]] const PreparedBarotropicClosureOperator &GetBarotropicClosureOperator() const noexcept;
        [[nodiscard]] const context::barotropic::BarotropicClosureLinearizationContext &
        GetBarotropicClosureContext() const noexcept;
        [[nodiscard]] const PreparedHydrostaticEquilibriumOperator &GetHydrostaticOperator() const noexcept;
        [[nodiscard]] const PreparedDisplacementResidualOperator &GetDisplacementOperator() const noexcept;
        [[nodiscard]] const PreparedMassNormalizationOperator &GetMassNormalizationOperator() const noexcept;
        [[nodiscard]] double ApplyDensityVolumeIntegralDensityAction(const mfem::Vector &densityDirection) const;
        [[nodiscard]] double
        ApplyDensityVolumeIntegralSurfaceShapeAction(const mfem::Vector &surfaceShapeDirection) const;
        [[nodiscard]] const PreparedPressureSurfaceConstraint &GetSurfaceConstraintOperator() const noexcept;
        [[nodiscard]] const deformation::PreparedDomainDeformationRuntime &GetDomainDeformation() const noexcept;
        [[nodiscard]] const mfem::Vector &GetSurfaceDeformationParameters() const;
        [[nodiscard]] const mfem::Vector &GetGeneratedVolumeDisplacement() const;
        [[nodiscard]] const mfem::Vector &GetFullMechanicalResidual() const;
        [[nodiscard]] const StellarEquilibriumDependencyStamp &GetGeneratedDisplacementDependency() const;

    private:
        struct ConstructionData;

        static ConstructionData MakeConstructionData(
            fem::FEM &f,
            deformation::PreparedDomainDeformationRuntime domainDeformation
        );

        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            models::CompiledFixedMass fixedMassConstraint,
            PressureSurfaceConstraintView surfaceConstraint,
            ConstructionData constructionData
        );

        void AssembleResidual();
        void VerifyPrepared() const;

        StellarEquilibriumRootManifest m_rootManifest;
        MPI_Comm m_communicator{MPI_COMM_NULL};
        mfem::Array<int> m_gravityStateOffsets;

        context::gravity_field::GravityFieldLinearizationContext m_gravityContext;
        GravityFieldJacobianOperator m_gravityJacobianOperator;
        GravityFieldOperator m_gravityOperator;

        PreparedBarotropicClosureOperator m_barotropicClosureOperator;
        PreparedHydrostaticEquilibriumOperator m_hydrostaticOperator;
        PreparedDisplacementResidualOperator m_displacementOperator;
        PreparedMassNormalizationOperator m_massNormalizationOperator;
        PreparedPressureSurfaceConstraint m_surfaceConstraintOperator;
        deformation::PreparedDomainDeformationRuntime m_domainDeformation;

        StellarEquilibriumDependencies m_preparedDependencies;
        StellarEquilibriumDependencyStamp m_generatedDisplacementDependency;
        deformation::DomainDeformationGeometryReport m_generatedGeometryReport;
        mfem::Vector m_surfaceDeformationParameters;
        mfem::Vector m_generatedVolumeDisplacement;
        mfem::Vector m_fullMechanicalResidual;
        mfem::Vector m_cachedResidual;
        models::CompiledFixedMass m_fixedMassConstraint;

        mutable PreparedStellarEquilibriumStatistics m_statistics;
        bool m_isPrepared{false};

        mfem::Vector m_gravityState;

        mutable mfem::Vector m_gravityDirection;
        mutable mfem::Vector m_volumeDisplacementDirection;
        mutable mfem::Vector m_fullMechanicalAction;
        mutable mfem::Vector m_surfaceShapeAction;
        mutable mfem::Vector m_pullbackDerivativeAction;
        mutable mfem::Vector m_densityVolumeIntegralAction;
    };
} // namespace mean_field::operators