module;

#include <compare>
#include <cstdint>

#include <mfem.hpp>

export module mean_field:operators.context.barotropic_closure_linearization;

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

export namespace mean_field::operators::context::barotropic {
    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 DensityDependencyTag { };
    struct EnthalpyDependencyTag { };
    struct DisplacementDependencyTag { };

    using DiscretizationDependency = DependencyStamp<DiscretizationDependencyTag>;
    using DensityDependency        = DependencyStamp<DensityDependencyTag>;
    using EnthalpyDependency       = DependencyStamp<EnthalpyDependencyTag>;
    using DisplacementDependency   = DependencyStamp<DisplacementDependencyTag>;

    struct BarotropicClosureDependencies final {
        DiscretizationDependency discretization;
        DensityDependency density;
        EnthalpyDependency enthalpy;
        DisplacementDependency displacement;

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

    struct BarotropicClosureStateView final {
        const mfem::Vector &density;
        const mfem::Vector &enthalpy;
        const mfem::Vector &displacement;
    };

    struct BarotropicClosurePreparationReport final {
        bool preparedStaticDependencies{false};
        bool preparedGeometryState{false};
        bool preparedBaseState{false};

        bool updatedDensity{false};
        bool updatedEnthalpy{false};
        bool updatedDisplacement{false};

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

    struct BarotropicClosurePreparationStatistics final {
        std::uint64_t staticPreparations{0};
        std::uint64_t geometryPreparations{0};
        std::uint64_t baseStatePreparations{0};

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

    class BarotropicClosureLinearizationContext final {
    public:
        BarotropicClosureLinearizationContext(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const field::FieldDofMap &densityMap,
            const field::FieldDofMap &enthalpyMap,
            const field::FieldDofMap &displacementMap
        );

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

        BarotropicClosurePreparationReport Prepare(
            const BarotropicClosureStateView &state,
            const BarotropicClosureDependencies &dependencies
        );

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] bool MatchesDependencies(const BarotropicClosureDependencies &dependencies) const noexcept;
        [[nodiscard]] const BarotropicClosureDependencies &GetDependencies() const;
        [[nodiscard]] const BarotropicClosurePreparationStatistics &GetPreparationStatistics() const noexcept;

        [[nodiscard]] const mfem::Vector &GetBaseDensity() const;
        [[nodiscard]] const mfem::Vector &GetBaseEnthalpy() const;
        [[nodiscard]] const mfem::Vector &GetDisplacement() const;

    private:
        void VerifyPrepared() const;

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

        int m_densitySize{0};
        int m_enthalpySize{0};
        int m_displacementSize{0};

        mfem::Vector m_baseDensity;
        mfem::Vector m_baseEnthalpy;
        mfem::Vector m_displacement;

        BarotropicClosureDependencies m_dependencies;
        BarotropicClosurePreparationStatistics m_statistics;
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators::context::barotropic