/*
 * Copyright (c) 2020-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/graph/format.hpp>

#include <cuda_runtime.h>

#include <raft/random/rng_state.hpp>

#include <rmm/device_uvector.hpp>

#include <cstdint>
#include <tuple>

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

/**
 * @brief Temporary enum covering RAFT generator types
 */
enum class GeneratorTypeT {
  kGenPC = 0,
  kGenPhilox,
};  // enum GeneratorTypeT

/**
 * @brief Different kinds of sampling algorithms
 */
enum class SamplingAlgoT {
  /** Reservoir sampling (Algo R) */
  kReservoirAlgoR,
  /** Reservoir sampling (Algo L, single-threaded) */
  kReservoirAlgoLST,
};  // enum SamplingAlgoT

/**
 * @brief Generate indexes given population sizes and a sample size,
 *        with or without replacement
 *
 * @param[out]   index          The (dense) index matrix. [on device]
 *                              [dim = `n_sizes x sample_size`]
 *                              In case `replace` is `false`, this may contain
 *                              `ops::gnn::graph::INVALID_ID<IdxT>`
 *                              if no index could be generated.
 * @param[inout] rng            RAFT RngState state object
 * @param[in]    sizes          Input array of population sizes [on device]
 *                              [len = `n_sizes`]
 * @param[in]    n_sizes        number of sizes to sample from.
 * @param[in]    sample_size    max number of indexes to be sampled per element
 *                              in `sizes`. Assumed to be <= 384 at the moment.
 * @param[in]    replace        If `true`, sample with replacement, otherwise
 *                              without replacement.
 * @param[in]    stream         cuda stream
 *
 @{
 */
void get_sampling_index(int32_t* index,
                        raft::random::RngState& rng,
                        const int32_t* sizes,
                        int32_t n_sizes,
                        int32_t sample_size,
                        bool replace,
                        cudaStream_t stream);
void get_sampling_index(int64_t* index,
                        raft::random::RngState& rng,
                        const int64_t* sizes,
                        int64_t n_sizes,
                        int32_t sample_size,
                        bool replace,
                        cudaStream_t stream);
/** @} */

/**
 * @brief Uniform graph neighboorhood sampling technique
 *
 * @note explicit self-loops in the input graph are always sampled
 *       (like any other edge)
 *
 * @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 = `n_nodes`].
 *                              If this is a `nullptr`, then all nodes
 *                              neighborhood will be sampled.
 * @param[in]    n_nodes        number of nodes to be sampled from.
 * @param[in]    sample_size    max number of nodes to be sampled per output node
 * @param[in]    type           sampling algorithm type
 * @param[in]    max_val        maximum node degree found in the graph. If not
 *                              used by the underlying algo, pass a `0`.
 * @param[in]    stream         cuda stream
 * @return tuple of device vectors representing CSR offsets and indices of
 *         the sub-sampled graph or COO source and destination indices
 *
 @{
 */
std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>> uniform_sample_coo(
  raft::random::RngState& rng,
  const fg_csr_s32_t& graph,
  const int32_t* nodes,
  int32_t n_nodes,
  int32_t sample_size,
  SamplingAlgoT type,
  int32_t max_val,
  cudaStream_t stream);
std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int64_t>> uniform_sample_coo(
  raft::random::RngState& rng,
  const fg_csr_s64_t& graph,
  const int64_t* nodes,
  int64_t n_nodes,
  int64_t sample_size,
  SamplingAlgoT type,
  int64_t max_val,
  cudaStream_t stream);
std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>> uniform_sample_csr(
  raft::random::RngState& rng,
  const fg_csr_s32_t& graph,
  const int32_t* nodes,
  int32_t n_nodes,
  int32_t sample_size,
  SamplingAlgoT type,
  int32_t max_val,
  cudaStream_t stream);
std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int64_t>> uniform_sample_csr(
  raft::random::RngState& rng,
  const fg_csr_s64_t& graph,
  const int64_t* nodes,
  int64_t n_nodes,
  int64_t sample_size,
  SamplingAlgoT type,
  int64_t max_val,
  cudaStream_t stream);
/** @} */

/**
 * @brief Uniform graph neighboorhood sampling technique
 *
 * @param[out]   mfg            The (ELLPACK/CSR) MFG. 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& mfg,
                    raft::random::RngState& rng,
                    const fg_csr_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& mfg,
                    raft::random::RngState& rng,
                    const fg_csr_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_csr_s32_t& mfg,
                    raft::random::RngState& rng,
                    const fg_csr_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_csr_s64_t& mfg,
                    raft::random::RngState& rng,
                    const fg_csr_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::gnn::graph
