/*
 * 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 {

/**
 * @brief Computes the forward pass for a vanilla aggregation (agg_simple)
 *        operating on a full graph (fg) in a node-to-node reduction (n2n)
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_fg_n2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* node_feat,
                           size_t dim_node,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
void agg_simple_fg_n2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* node_feat,
                           size_t dim_node,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
void agg_simple_fg_n2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* node_feat,
                           size_t dim_node,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
void agg_simple_fg_n2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* node_feat,
                           size_t dim_node,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a vanilla aggregation (agg_simple)
 *        operating on a full graph (fg) in a node-to-node reduction (n2n)
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `.n_nodes x dim_node`]
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `.n_nodes x dim_node`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_fg_n2n_bwd(float* d_node_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
void agg_simple_fg_n2n_bwd(float* d_node_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_node,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
void agg_simple_fg_n2n_bwd(float* d_node_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
void agg_simple_fg_n2n_bwd(float* d_node_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_node,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           bool ignore_self,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple aggregation (agg_simple)
 *        using edge features in an edge-to-node reduction (e2n)
 *        operating on a full graph (fg)
 *
 * @note The `eidx` version uses an edge index to obtain the index of an edge
 * feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary. In many cases including if there is no edge feature index,
 * the number of rows is `fg.n_indices` for a single graph or
 * `fg_batch.n_edges` for a batch of graphs.
 * Below we'll refer to the number of rows simply as `n_edges`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x dim_edge`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_nodes x dim_edge`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  edge_feat       the input edge features. [on device]
 *                             [dim = `.n_edges x dim_edge`]
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_fg_e2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* edge_feat,
                           size_t dim_edge,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_fwd(float* out,
                           int32_t* out_pos,
                           const float* edge_feat,
                           size_t dim_edge,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* edge_feat,
                           size_t dim_edge,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_fwd(float* out,
                           int64_t* out_pos,
                           const float* edge_feat,
                           size_t dim_edge,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_fwd(float* out,
                                int32_t* out_pos,
                                const float* edge_feat,
                                size_t dim_edge,
                                const fg_csr_s32_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_fwd(float* out,
                                int32_t* out_pos,
                                const float* edge_feat,
                                size_t dim_edge,
                                const fg_csr_batch_s32_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_fwd(float* out,
                                int64_t* out_pos,
                                const float* edge_feat,
                                size_t dim_edge,
                                const fg_csr_s64_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_fwd(float* out,
                                int64_t* out_pos,
                                const float* edge_feat,
                                size_t dim_edge,
                                const fg_csr_batch_s64_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a simple aggregation (agg_simple)
 *        using edge features in an edge-to-node reduction (e2n)
 *        operating on a full graph (fg)
 *
 * @note The `eidx` version uses an edge index to obtain the index of an edge
 * feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `d_edge_feat` is
 * arbitrary. In many cases including if there is no edge feature index,
 * the number of rows is `fg.n_indices` for a single graph or
 * `fg_batch.n_edges` for a batch of graphs.
 * Below we'll refer to the number of rows simply as `n_edges`.
 *
 * @param[out] d_edge_feat     the output gradients of input edge features
 *                             [on device] [dim = `n_edges x dim_edge`]
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `fg.n_nodes x dim_edge`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_nodes x dim_edge`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_fg_e2n_bwd(float* d_edge_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           const fg_csr_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_bwd(float* d_edge_feat,
                           const float* dout,
                           const int32_t* out_pos,
                           size_t dim_edge,
                           const fg_csr_batch_s32_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_bwd(float* d_edge_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           const fg_csr_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_bwd(float* d_edge_feat,
                           const float* dout,
                           const int64_t* out_pos,
                           size_t dim_edge,
                           const fg_csr_batch_s64_t& fg,
                           AggOpT op,
                           const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_bwd(float* d_edge_feat,
                                const float* dout,
                                const int32_t* out_pos,
                                size_t dim_edge,
                                const fg_csr_s32_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_bwd(float* d_edge_feat,
                                const float* dout,
                                const int32_t* out_pos,
                                size_t dim_edge,
                                const fg_csr_batch_s32_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_bwd(float* d_edge_feat,
                                const float* dout,
                                const int64_t* out_pos,
                                size_t dim_edge,
                                const fg_csr_s64_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_simple_fg_e2n_eidx_bwd(float* d_edge_feat,
                                const float* dout,
                                const int64_t* out_pos,
                                size_t dim_edge,
                                const fg_csr_batch_s64_t& fg,
                                AggOpT op,
                                const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple aggregation (agg_simple)
 *        using edge features and node features in an edge-to-node reduction (e2n)
 *        and a node-to-node reduction (n2n) operating on a full graph (fg).
 *        The results of both individual aggregations are simply concatenated.
 *
 * @note The `eidx` version uses an edge index to obtain the index of an edge
 * feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary. In many cases including if there is no edge feature index,
 * the number of rows is `fg.n_indices` for a single graph or
 * `fg_batch.n_edges` for a batch of graphs.
 * Below we'll refer to the number of rows simply as `n_edges`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_edge)`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_edge)`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 * @param[in]  edge_feat       the input edge features. [on device]
 *                             [dim = `.n_edges x dim_edge`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  ignore_self     ignore the self node in the aggregation of node features
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_fg_n2n_e2n_fwd(float* out,
                               int32_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_fwd(float* out,
                               int32_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_fwd(float* out,
                               int64_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_fwd(float* out,
                               int64_t* out_pos,
                               const float* node_feat,
                               const float* edge_feat,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_fwd(float* out,
                                    int32_t* out_pos,
                                    const float* node_feat,
                                    const float* edge_feat,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_s32_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_fwd(float* out,
                                    int32_t* out_pos,
                                    const float* node_feat,
                                    const float* edge_feat,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_batch_s32_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_fwd(float* out,
                                    int64_t* out_pos,
                                    const float* node_feat,
                                    const float* edge_feat,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_s64_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_fwd(float* out,
                                    int64_t* out_pos,
                                    const float* node_feat,
                                    const float* edge_feat,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_batch_s64_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a simple aggregation (agg_simple)
 *        using edge features and node features in an edge-to-node reduction (e2n)
 *        and a node-to-node reduction (n2n) operating on a full graph (fg).
 *        The results of both individual aggregations are simply concatenated.
 *
 * @note The `eidx` version uses an edge index to obtain the index of an edge
 * feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `d_edge_feat` is
 * arbitrary. In many cases including if there is no edge feature index,
 * the number of rows is `fg.n_indices` for a single graph or
 * `fg_batch.n_edges` for a batch of graphs.
 * Below we'll refer to the number of rows simply as `n_edges`.
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `fg.n_nodes x dim_node`]
 * @param[out] d_edge_feat     the output gradients of input edge features
 *                             [on device] [dim = `n_edges x dim_edge`]
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_edge)`]
 * @param[in]  out_pos         the position of min/max values [on device]
 *                             [dim = `fg.n_nodes x (dim_node + dim_edge)`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  fg              the input CSR full graph (or batch of graphs)
 * @param[in]  op              aggregation operation
 * @param[in]  ignore_self     ignore the self node in the aggregation of node features
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               const float* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_s32_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               const float* dout,
                               const int32_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_batch_s32_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               const float* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_s64_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_bwd(float* d_node_feat,
                               float* d_edge_feat,
                               const float* dout,
                               const int64_t* out_pos,
                               size_t dim_node,
                               size_t dim_edge,
                               const fg_csr_batch_s64_t& fg,
                               AggOpT op,
                               bool ignore_self,
                               const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_bwd(float* d_node_feat,
                                    float* d_edge_feat,
                                    const float* dout,
                                    const int32_t* out_pos,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_s32_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_bwd(float* d_node_feat,
                                    float* d_edge_feat,
                                    const float* dout,
                                    const int32_t* out_pos,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_batch_s32_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_bwd(float* d_node_feat,
                                    float* d_edge_feat,
                                    const float* dout,
                                    const int64_t* out_pos,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_s64_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_simple_fg_n2n_e2n_eidx_bwd(float* d_node_feat,
                                    float* d_edge_feat,
                                    const float* dout,
                                    const int64_t* out_pos,
                                    size_t dim_node,
                                    size_t dim_edge,
                                    const fg_csr_batch_s64_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
