/*
 * Copyright (c) 2020-2023, 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/graph/sampling.hpp>
#include <cugraph-ops/test/graph/format_mfg.hpp>
#include <cugraph-ops/test/graph/sampling_mfg.hpp>

#include <cuda_runtime.h>

#include <raft/random/rng_state.hpp>

#include <rmm/device_uvector.hpp>

#include <cstdint>
#include <tuple>

namespace cugraph::ops::graph {

/**
 * @brief Uniform graph neighboorhood sampling technique
 *
 * @param[out]   sub_graph      The sampled sub-graph in MFG format (ELLPACK/CSR).
                                The `neighbors` matrix
 *                              and `neighbor_counts` array will be filled.
 *                              For CSR, `offsets`, `indices` and `n_indices`
 *                              will be filled.
 * @param[inout] rng            RAFT RngState state object
 * @param[in]    graph          input graph
 * @param[in]    nodes          list of node indices whose neighbors need to be
 *                              sampled [on device] [len = `mfg.n_out_nodes`].
 *                              If this is a `nullptr`, then all nodes
 *                              neighborhood will be sampled.
 * @param[in]    type           sampling type
 * @param[in]    max_val        maximum node degree found in the graph. If not
 *                              used by the underlying algo, pass a `0`.
 * @param[in]    sample_self    if `true`, explicit self loops in the graph can
 *                              be sampled as if they were regular neighbors.
 * @param[inout] n_indices      Intermediate buffer for the number of edges.
 *                              Unused in the ELLPACK variant [on host, pinned]
 *                              [len = `1`]
 * @param[inout] workspace      scratch buffer. [on device]
 *                              [len = `workspace_size`] Pass a `nullptr` in
 *                              order to get to know its size.
 * @param[inout] workspace_size workspace size in bytes. If `workspace` is a
 *                              `nullptr`, then this will be computed and the
 *                              caller is expected to allocate the workspace
 *                              buffer of this size.
 * @param[in]    stream         cuda stream
 *
 * @note It does not include the current node in the output neighborhood.
 *
 @{
 */
void uniform_sample(mfg_ellpack_s32_t& sub_graph,
                    raft::random::RngState& rng,
                    const csc_s32_t& graph,
                    const int32_t* nodes,
                    SamplingAlgoT type,
                    int32_t max_val,
                    bool sample_self,
                    int32_t* n_indices,
                    void* workspace,
                    size_t& workspace_size,
                    cudaStream_t stream);
void uniform_sample(mfg_ellpack_s64_t& sub_graph,
                    raft::random::RngState& rng,
                    const csc_s64_t& graph,
                    const int64_t* nodes,
                    SamplingAlgoT type,
                    int64_t max_val,
                    bool sample_self,
                    int64_t* n_indices,
                    void* workspace,
                    size_t& workspace_size,
                    cudaStream_t stream);
void uniform_sample(mfg_csc_s32_t& sub_graph,
                    raft::random::RngState& rng,
                    const csc_s32_t& graph,
                    const int32_t* nodes,
                    SamplingAlgoT type,
                    int32_t max_val,
                    bool sample_self,
                    int32_t* n_indices,
                    void* workspace,
                    size_t& workspace_size,
                    cudaStream_t stream);
void uniform_sample(mfg_csc_s64_t& sub_graph,
                    raft::random::RngState& rng,
                    const csc_s64_t& graph,
                    const int64_t* nodes,
                    SamplingAlgoT type,
                    int64_t max_val,
                    bool sample_self,
                    int64_t* n_indices,
                    void* workspace,
                    size_t& workspace_size,
                    cudaStream_t stream);
/** @} */

}  // namespace cugraph::ops::graph
