/*
 * 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/cuda/stream.hpp>
#include <cugraph-ops/graph/format.hpp>
#include <cugraph-ops/operators/common.hpp>

#include <cstdint>

namespace cugraph::ops::gnn {

// in case we expose functions for aggregating node (and node + edge) features
// note that both node and edge features are assumed to be in a contiguous
// matrix, and that the graph either does not have edge feature indices or
// the indices are simply a permutation of the edge IDs

/**
 * @brief Computes the forward pass for a pooling operation (pool_fg) aggregating
 *        node features from an entire graph of batch of graphs in a node-to-scalar
 *        reduction (n2s)
 *
 * @note For these operations, we assume that node features of each
 *       graph in the batch (or of a single graph) are in a contiguous matrix.
 *
 * @param[out]   out           the output aggregated embeddings. [on device]
 *                             [dim = `n_graphs x out_size`]
 * @param[out]   out_pos       location of the max/min embeddings. This will be
 *                             written to only for kMin/kMax aggregations and
 *                             when this pointer is NOT a `nullptr`. [on device]
 *                             [dim = `n_graphs x out_size`]
 * @param[in]    node_feat     the input node features. [on device]
 *                             [dim = `fg.n_nodes x dim`]
 * @param[in]    dim           dimensionality of the node features
 * @param[in]    fg            the input CSR full graph (or batch of graphs)
 * @param[in]    op            aggregation operation
 * @param[in]    use_atomics   If set to `true` and op is sum or mean, the
 *                             operation may not be deterministic but does not
 *                             require any workspace.
 *                             In this case (only if op is sum or mean),
 *                             `workspace` can be passed as `nullptr` and
 *                             `w_size` as `0` to immediately compute the output
 *                             without allocating an additional buffer.
 * @param[inout] workspace     workspace required for computation.
 *                             May be `nullptr` if we require no workspace.
 * @param[inout] w_size        workspace size required (denotes number of
 *                             elements expected for `workspace`).
 *                             If the input value is negative, compute the
 *                             workspace size and return.
 *                             The size depends on `dim` and batch size
 *                             (if we have a batch of graphs).
 *                             The returned value is always non-negative.
 * @param[in]    stream        cuda stream
 * @{
 */
void pool_fg_n2s_fwd(float* out,
                     int32_t* out_pos,
                     const float* node_feat,
                     size_t dim,
                     const fg_csr_s32_t& fg,
                     AggOpT op,
                     bool use_atomics,
                     int64_t* workspace,
                     int64_t& w_size,
                     const cuda::stream& stream);
void pool_fg_n2s_fwd(float* out,
                     int32_t* out_pos,
                     const float* node_feat,
                     size_t dim,
                     const fg_csr_batch_s32_t& fg,
                     AggOpT op,
                     bool use_atomics,
                     int64_t* workspace,
                     int64_t& w_size,
                     const cuda::stream& stream);
void pool_fg_n2s_fwd(float* out,
                     int64_t* out_pos,
                     const float* node_feat,
                     size_t dim,
                     const fg_csr_s64_t& fg,
                     AggOpT op,
                     bool use_atomics,
                     int64_t* workspace,
                     int64_t& w_size,
                     const cuda::stream& stream);
void pool_fg_n2s_fwd(float* out,
                     int64_t* out_pos,
                     const float* node_feat,
                     size_t dim,
                     const fg_csr_batch_s64_t& fg,
                     AggOpT op,
                     bool use_atomics,
                     int64_t* workspace,
                     int64_t& w_size,
                     const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a pooling operation (pool_fg) aggregating
 *        node features from an entire graph of batch of graphs in a node-to-scalar
 *        reduction (n2s)
 *
 * @note For these operations, we assume that node features of each
 *       graph in the batch (or of a single graph) are in a contiguous matrix.
 *
 * @param[out] d_node_feat     the output node feature gradients. [on device]
 *                             [dim = `fg.n_nodes x dim`]
 * @param[in]  d_out           the gradient on output features. [on device]
 *                             [dim = `n_graphs x out_size`]
 * @param[in]  out_pos         location of the max/min embeddings from forward.
 *                             Must be the same pointer as in forward.
 *                             [on device] [dim = `n_graphs x out_size`]
 * @param[in]  dim             dimensionality of the node features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void pool_fg_n2s_bwd(float* d_node_feat,
                     const float* d_out,
                     const int32_t* out_pos,
                     size_t dim,
                     const fg_csr_s32_t& fg,
                     AggOpT op,
                     const cuda::stream& stream);
void pool_fg_n2s_bwd(float* d_node_feat,
                     const float* d_out,
                     const int32_t* out_pos,
                     size_t dim,
                     const fg_csr_batch_s32_t& fg,
                     AggOpT op,
                     const cuda::stream& stream);
void pool_fg_n2s_bwd(float* d_node_feat,
                     const float* d_out,
                     const int64_t* out_pos,
                     size_t dim,
                     const fg_csr_s64_t& fg,
                     AggOpT op,
                     const cuda::stream& stream);
void pool_fg_n2s_bwd(float* d_node_feat,
                     const float* d_out,
                     const int64_t* out_pos,
                     size_t dim,
                     const fg_csr_batch_s64_t& fg,
                     AggOpT op,
                     const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
