module;

#include <compare>
#include <cstdint>
#include <optional>

#include <mfem.hpp>

export module mean_field:operators.prepared_rotational_displacement_force;

export import :fem;
export import :mapping.domain_mapper;
export import :operators.context.rotational_displacement_force;
export import :physics.rigid_rotation;
export import :utils.blocks;

export namespace mean_field::operators {
    struct PreparedRotationalDisplacementForceReport final {
        context::rotational_displacement_force::RotationalDisplacementForcePreparationReport contextReport;

        bool updatedRotation{false};
        bool preparedResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return contextReport.DidAnyWork() || updatedRotation || preparedResidual;
        }
    };

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

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

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

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

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

        PreparedRotationalDisplacementForceOperator(const PreparedRotationalDisplacementForceOperator &) = delete;

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

        PreparedRotationalDisplacementForceOperator(PreparedRotationalDisplacementForceOperator &&)            = delete;

        PreparedRotationalDisplacementForceOperator &operator=(PreparedRotationalDisplacementForceOperator &&) = delete;

        PreparedRotationalDisplacementForceReport Prepare(
            const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
            const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
            const physics::RigidRotation &rotation
        );

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

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

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

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

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

        [[nodiscard]] const context::rotational_displacement_force::RotationalDisplacementForcePreparationStatistics &
        GetContextPreparationStatistics() const noexcept;

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

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

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

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

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

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

        [[nodiscard]] const context::rotational_displacement_force::RotationalDisplacementForceLinearizationContext &
        GetContext() const noexcept;

    private:
        void VerifyPrepared() const;

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

        context::rotational_displacement_force::RotationalDisplacementForceLinearizationContext m_context;

        std::optional<physics::RigidRotation> m_rotation;
        mfem::Vector m_cachedResidual;

        context::rotational_displacement_force::RotationalDisplacementForceDependencies m_preparedDependencies;

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

        mutable PreparedRotationalDisplacementForceColumnStatistics m_densityJacobianStatistics;

        mutable PreparedRotationalDisplacementForceColumnStatistics m_displacementJacobianStatistics;

        mutable PreparedRotationalDisplacementForceCompleteStatistics m_completeJacobianStatistics;

        bool m_isPrepared{false};
    };

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

    class PreparedRotationalDisplacementForceJacobianOperator final : public mfem::Operator {
    public:
        PreparedRotationalDisplacementForceJacobianOperator(
            const RotationalDisplacementForceLayout &layout,
            const PreparedRotationalDisplacementForceOperator &preparedOperator
        );

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

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

    private:
        RotationalDisplacementForceLayout m_layout;
        const PreparedRotationalDisplacementForceOperator &m_preparedOperator;
    };
} // namespace mean_field::operators