/*
 * Copyright (c) 2021-2022, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/graph/format.hpp>

#include <raft/random/rng_state.hpp>

namespace cugraph::ops::gnn::graph {

/**
 * @brief Generate a random graph in CSR format on device
 *
 * @param[inout] graph          graph to be generated
 * @param[inout] r              RAFT RngState object
 * @param[in]    nodes          number of nodes
 * @param[in]    nnz            number of indexes (neighbors)
 * @param[in]    is_undirected  if `true`, generate an undirected graph.
 *                              In this case, the actual number of indices is:
 *                              `2 * floor(nnz / 2)`
 * @param[in]    stream         cuda stream
 *
 * @note this could generate duplicate edges (multi-graphs)!!!
 *
 * @{
 */
void generate_csr_on_device(fg_csr_s32_t& graph,
                            raft::random::RngState& r,
                            int32_t nodes,
                            int32_t nnz,
                            bool is_undirected,
                            cuda::stream& stream);
void generate_csr_on_device(fg_csr_s64_t& graph,
                            raft::random::RngState& r,
                            int64_t nodes,
                            int64_t nnz,
                            bool is_undirected,
                            cuda::stream& stream);
/** @} */

/**
 * @brief Delete the previously generated CSR graph from device
 *
 * @param[inout] graph  graph to be deleted
 * @param[in]    stream cuda stream
 *
 * @{
 */
void delete_csr_from_device(fg_csr_s32_t& graph, cuda::stream& stream);
void delete_csr_from_device(fg_csr_s64_t& graph, cuda::stream& stream);
/** @} */

/**
 * @brief Generate a random set of graphs in the sequence of CSR format on device
 *
 * @param[inout] graph          graphs to be generated
 * @param[inout] r              RAFT RngState object
 * @param[in]    n_graphs       number of graphs to be generated
 * @param[in]    nodes          number of nodes. It is currently assumed that every graph
 *                              has the same number of nodes
 * @param[in]    nnz            number of indexes (neighbors). It is currently assumed that
 *                              every graph has the same number of edges
 * @param[in]    is_undirected  if `true`, generate an undirected graph.
 *                              In this case, the actual number of indices per graph:
 *                              `2 * floor(nnz / 2)`
 * @param[in]    stream         cuda stream
 *
 * @note this could generate duplicate edges (multi-graphs)!!!
 *
 * @{
 */
void generate_csr_seq_on_device(fg_csr_seq_s32_t& graph,
                                raft::random::RngState& r,
                                int32_t n_graphs,
                                int32_t nodes,
                                int32_t nnz,
                                bool is_undirected,
                                cuda::stream& stream);
void generate_csr_seq_on_device(fg_csr_seq_s64_t& graph,
                                raft::random::RngState& r,
                                int64_t n_graphs,
                                int64_t nodes,
                                int64_t nnz,
                                bool is_undirected,
                                cuda::stream& stream);
/** @} */

/**
 * @brief Delete the previously generated sequence of CSR graphs from device
 *
 * @param[inout] graph  graphs to be deleted
 * @param[in]    stream cuda stream
 *
 * @{
 */
void delete_csr_seq_from_device(fg_csr_seq_s32_t& graph, cuda::stream& stream);
void delete_csr_seq_from_device(fg_csr_seq_s64_t& graph, cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn::graph
