module;

#include <compare>
#include <cstdint>

#include <mfem.hpp>

export module mean_field:operators.context.hydrostatic_equilibrium;

export import :fem;
export import :mapping.domain_mapper;

export namespace mean_field::operators::context::hydrostatic {
    template <typename Tag> struct DependencyStamp {
        std::uint64_t identity{0};
        std::uint64_t revision{0};

        [[nodiscard]] constexpr bool
        CanFollow(const DependencyStamp &prepared) const noexcept {
            return identity != prepared.identity ||
                   revision >= prepared.revision;
        }

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

    struct DiscretizationDependencyTag { };
    struct EnthalpyDependencyTag { };
    struct GravityPotentialDependencyTag { };
    struct DisplacementDependencyTag { };
    struct RotationDependencyTag { };
    struct BernoulliConstantDependencyTag { };

    using DiscretizationDependency =
        DependencyStamp<DiscretizationDependencyTag>;

    using EnthalpyDependency = DependencyStamp<EnthalpyDependencyTag>;

    using GravityPotentialDependency =
        DependencyStamp<GravityPotentialDependencyTag>;

    using DisplacementDependency = DependencyStamp<DisplacementDependencyTag>;

    using RotationDependency     = DependencyStamp<RotationDependencyTag>;

    using BernoulliConstantDependency =
        DependencyStamp<BernoulliConstantDependencyTag>;

    struct HydrostaticEquilibriumDependencies {
        DiscretizationDependency discretization;
        EnthalpyDependency enthalpy;
        GravityPotentialDependency gravityPotential;
        DisplacementDependency displacement;
        RotationDependency rotation;
        BernoulliConstantDependency bernoulliConstant;

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

    struct HydrostaticEquilibriumStateView {
        const mfem::Vector &enthalpy;
        const mfem::Vector &gravityPotential;
        const mfem::Vector &displacement;
        double bernoulliConstant{0.0};
    };

    struct HydrostaticPreparationReport {
        bool preparedStaticDependencies{false};
        bool preparedGeometryState{false};
        bool preparedRotationDependencies{false};
        bool preparedBaseState{false};

        bool updatedEnthalpy{false};
        bool updatedGravityPotential{false};
        bool updatedDisplacement{false};
        bool updatedBernoulliConstant{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return preparedStaticDependencies || preparedGeometryState ||
                   preparedRotationDependencies || preparedBaseState;
        }
    };

    struct HydrostaticPreparationStatistics {
        std::uint64_t staticPreparations{0};
        std::uint64_t geometryPreparations{0};
        std::uint64_t rotationPreparations{0};
        std::uint64_t baseStatePreparations{0};

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

    class HydrostaticEquilibriumContext {
    public:
        HydrostaticEquilibriumContext(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper
        );

        HydrostaticEquilibriumContext(const HydrostaticEquilibriumContext &) =
            delete;

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

        HydrostaticEquilibriumContext(HydrostaticEquilibriumContext &&) =
            delete;

        HydrostaticEquilibriumContext &
        operator=(HydrostaticEquilibriumContext &&) = delete;

        HydrostaticPreparationReport Prepare(
            const HydrostaticEquilibriumStateView &state,
            const HydrostaticEquilibriumDependencies &dependencies
        );

        [[nodiscard]] bool IsPrepared() const noexcept;

        [[nodiscard]] bool MatchesDependencies(
            const HydrostaticEquilibriumDependencies &dependencies
        ) const noexcept;

        [[nodiscard]] const HydrostaticEquilibriumDependencies &
        GetDependencies() const;

        [[nodiscard]] const HydrostaticPreparationStatistics &
        GetPreparationStatistics() const noexcept;

        [[nodiscard]] const mfem::Vector &GetBaseEnthalpyTrue() const;

        [[nodiscard]] const mfem::Vector &GetBaseGravityPotentialTrue() const;

        [[nodiscard]] const mfem::Vector &GetDisplacementTrue() const;

        [[nodiscard]] double GetBernoulliConstant() const;

    private:
        void VerifyPrepared() const;

        const fem::FEM &m_f;
        const mapping::DomainMapperStateless &m_domainMapper;

        mfem::Vector m_baseEnthalpyTrue;
        mfem::Vector m_baseGravityPotentialTrue;
        mfem::Vector m_displacementTrue;
        double m_bernoulliConstant{0.0};

        HydrostaticEquilibriumDependencies m_dependencies;
        HydrostaticPreparationStatistics m_statistics;
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators::context::hydrostatic