module;

#include <compare>
#include <cstdint>

#include <mfem.hpp>

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 :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 :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 displacement;
        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;
        bool assembledResidual{false};

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

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

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

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

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

    class PreparedStellarEquilibriumOperator final : public mfem::Operator {
    public:
        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            double targetMass
        );

        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            const models::StellarModel &stellarModel
        );

        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
        );

        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 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;

    private:
        struct ConstructionData;

        static ConstructionData MakeConstructionData(fem::FEM &f);

        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            double targetMass,
            ConstructionData constructionData
        );

        void AssembleResidual();
        void VerifyPrepared() const;

        StellarEquilibriumLayout m_layout;
        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;

        StellarEquilibriumDependencies m_preparedDependencies;
        mfem::Vector m_cachedResidual;
        double m_targetMass{0.0};

        mutable PreparedStellarEquilibriumStatistics m_statistics;
        bool m_isPrepared{false};

        mfem::Vector m_gravityState;

        mutable mfem::Vector m_gravityDirection;
    };
} // namespace mean_field::operators