module;

#include <compare>
#include <cstdint>

#include <mfem.hpp>

export module mean_field:operators.prepared_gravity_displacement_force;

export import :fem;
export import :mapping.domain_mapper;
export import :operators.context.gravity_field;
export import :utils.blocks;

export namespace mean_field::operators {
    struct PreparedGravityDisplacementForceReport final {
        bool preparedResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return preparedResidual;
        }
    };

    struct PreparedGravityDisplacementForceColumnStatistics final {
        std::uint64_t applications{0};

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

    struct PreparedGravityDisplacementForceCompleteStatistics final {
        std::uint64_t applications{0};

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

    class PreparedGravityDisplacementForceOperator final {
    public:
        PreparedGravityDisplacementForceOperator(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const context::gravity_field::GravityFieldLinearizationContext &gravityContext
        );

        PreparedGravityDisplacementForceOperator(const PreparedGravityDisplacementForceOperator &)            = delete;

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

        PreparedGravityDisplacementForceOperator(PreparedGravityDisplacementForceOperator &&)                 = delete;

        PreparedGravityDisplacementForceOperator &operator=(PreparedGravityDisplacementForceOperator &&)      = delete;

        /*
         * Freeze the already-prepared shared gravity context. The caller must
         * first prepare GravityFieldLinearizationContext for the desired
         * state. Repeated calls with unchanged revisions do no work.
         */
        PreparedGravityDisplacementForceReport Prepare();

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

        void ApplyDensityJacobianAction(
            const mfem::Vector &densityVariation,
            mfem::Vector &action
        ) const;

        void ApplyGravityGradientJacobianAction(
            const mfem::Vector &gravityGradientVariation,
            mfem::Vector &action
        ) const;

        void ApplyDisplacementJacobianAction(
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

        void ApplyCompleteJacobianAction(
            const mfem::Vector &densityVariation,
            const mfem::Vector &displacementVariation,
            const mfem::Vector &gravityGradientVariation,
            mfem::Vector &action
        ) const;

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

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

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

        [[nodiscard]] const PreparedGravityDisplacementForceColumnStatistics &
        GetDensityJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceColumnStatistics &
        GetGravityGradientJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceColumnStatistics &
        GetDisplacementJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceCompleteStatistics &
        GetCompleteJacobianStatistics() const noexcept;

        [[nodiscard]] const fem::FEM &GetFEM() const noexcept;

        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetGravityContext() const noexcept;

    private:
        void VerifyPrepared() const;

        const fem::FEM &m_fem;
        const mapping::DomainMapperStateless &m_domainMapper;
        const context::gravity_field::GravityFieldLinearizationContext &m_gravityContext;

        context::gravity_field::GravityFieldRevisions m_preparedRevisions;
        mfem::Vector m_cachedResidual;

        std::uint64_t m_residualPreparationCount{0};
        mutable std::uint64_t m_residualApplicationCount{0};

        mutable PreparedGravityDisplacementForceColumnStatistics m_densityJacobianStatistics;

        mutable PreparedGravityDisplacementForceColumnStatistics m_gravityGradientJacobianStatistics;

        mutable PreparedGravityDisplacementForceColumnStatistics m_displacementJacobianStatistics;

        mutable PreparedGravityDisplacementForceCompleteStatistics m_completeJacobianStatistics;

        bool m_isPrepared{false};
    };

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

    class PreparedGravityDisplacementForceJacobianOperator final : public mfem::Operator {
    public:
        PreparedGravityDisplacementForceJacobianOperator(
            const GravityDisplacementForceLayout &layout,
            const PreparedGravityDisplacementForceOperator &preparedOperator
        );

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override;

        [[nodiscard]] const GravityDisplacementForceLayout &GetLayout() const noexcept;

    private:
        GravityDisplacementForceLayout m_layout;
        const PreparedGravityDisplacementForceOperator &m_preparedOperator;
    };
} // namespace mean_field::operators