module;

#include <cstdint>

#include <mfem.hpp>

export module mean_field:operators.context.barotropic_closure_linearization;

export import :fem;
export import :mapping.domain_mapper;
export import :operators.prepared_barotropic_closure;
export import :physics.barotrope;

export namespace mean_field::operators::context::barotropic {
    struct BarotropicClosureRevisions final {
        std::uint64_t density      = 0;
        std::uint64_t enthalpy     = 0;
        std::uint64_t displacement = 0;

        [[nodiscard]] bool
        operator==(const BarotropicClosureRevisions &) const noexcept = default;
    };

    class BarotropicClosureLinearizationContext final {
    public:
        BarotropicClosureLinearizationContext(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const physics::PolytropicBarotrope &barotrope
        );

        void Prepare(
            const mfem::Vector &baseDensityTrue,
            const mfem::Vector &baseEnthalpyTrue,
            const mfem::Vector &displacementTrue,
            const BarotropicClosureRevisions &revisions
        );

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

        [[nodiscard]] bool MatchesRevisions(
            const BarotropicClosureRevisions &revisions
        ) const noexcept;

        [[nodiscard]] std::uint64_t GetPreparationCount() const noexcept;

        [[nodiscard]] const BarotropicClosureRevisions &GetRevisions() const;

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

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

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

        [[nodiscard]]
        const PreparedBarotropicClosureOperator &GetOperator() const noexcept;

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

    private:
        void VerifyPrepared() const;

        const fem::FEM &m_f;

        PreparedBarotropicClosureOperator m_operator;

        mfem::Vector m_baseDensityTrue;
        mfem::Vector m_baseEnthalpyTrue;
        mfem::Vector m_displacementTrue;

        BarotropicClosureRevisions m_revisions;

        std::uint64_t m_preparationCount = 0;
        bool m_isPrepared                = false;
    };
} // namespace mean_field::operators::context::barotropic