module;

#include <compare>
#include <cstdint>

#include <mfem.hpp>

export module mean_field:operators.context.pressure_force;

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

export namespace mean_field::operators::context::pressure_force {
    template <typename Tag> struct DependencyStamp final {
        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 final { };
    struct EnthalpyDependencyTag final { };
    struct DisplacementDependencyTag final { };

    using DiscretizationDependency = DependencyStamp<DiscretizationDependencyTag>;

    using EnthalpyDependency       = DependencyStamp<EnthalpyDependencyTag>;

    using DisplacementDependency   = DependencyStamp<DisplacementDependencyTag>;

    struct PressureForceDependencies final {
        DiscretizationDependency discretization;
        EnthalpyDependency enthalpy;
        DisplacementDependency displacement;

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

    /*
     * Frozen solver-facing state.
     *
     * Both vectors use their registered FieldDof coordinates.
     *
     * Under the current registry:
     *
     *   enthalpy     -> Stellar -> reduced
     *   displacement -> All     -> identity/full
     */
    struct PressureForceStateView final {
        const mfem::Vector &enthalpy;
        const mfem::Vector &displacement;
    };

    struct PressureForcePreparationReport final {
        bool preparedStaticDependencies{false};
        bool preparedGeometryState{false};
        bool preparedMaterialState{false};

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

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

    struct PressureForcePreparationStatistics final {
        std::uint64_t staticPreparations{0};
        std::uint64_t geometryPreparations{0};
        std::uint64_t materialPreparations{0};

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

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

        PressureForceLinearizationContext(const PressureForceLinearizationContext &)            = delete;

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

        PressureForceLinearizationContext(PressureForceLinearizationContext &&)                 = delete;

        PressureForceLinearizationContext &operator=(PressureForceLinearizationContext &&)      = delete;

        PressureForcePreparationReport Prepare(
            const PressureForceStateView &state,
            const PressureForceDependencies &dependencies
        );

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

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

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

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

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

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

    private:
        void VerifyPrepared() const;

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

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

        PressureForcePreparationStatistics m_statistics;
        PressureForceDependencies m_dependencies;

        bool m_isPrepared{false};
    };
} // namespace mean_field::operators::context::pressure_force