/*
 * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#pragma once

#include <cuda_runtime.h>

#include <cstdint>

namespace cugraph::ops::graph {

/**
 * @brief Generates the next batch of nodes for training
 *
 * @param[out]   batch          output batch of nodes. [on device] [len = `size`]
 * @param[out]   next_start     state used for iteratively generating batch of nodes.
 *                              Before the start of the first batch, these are
 *                              expected to be initialized to zero [pinned] [len = 2]
 * @param[in]    size           batch size
 * @param[in]    n              total number of nodes
 * @param[in]    xor_key        random key that's expected to be the same for the
 *                              full epoch.
 * @param[in]    workspace      workspace for computation. [on device]
 *                              [len = `workspace_size`]. Pass a `nullptr` and
 *                              get the workspace size to be allocated.
 * @param[inout] workspace_size workspace size (in B)
 * @param[in]    stream         cuda stream
 * @{
 */
void get_next_output_batch(int32_t* batch,
                           int32_t* next_start,
                           int32_t size,
                           int32_t n,
                           uint64_t xor_key,
                           void* workspace,
                           size_t& workspace_size,
                           cudaStream_t stream);
void get_next_output_batch(int64_t* batch,
                           int64_t* next_start,
                           int64_t size,
                           int64_t n,
                           uint64_t xor_key,
                           void* workspace,
                           size_t& workspace_size,
                           cudaStream_t stream);
/** @} */

}  // namespace cugraph::ops::graph
