/*
 * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#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 {

/**
 * @brief Computes the forward pass for full-graph aggregation using node features
 *        while concatenting with the orignal node features, the actual aggregation
 *        is done weighting each neighbor according to its edge weight for each
 *        output node (agg_concat_weighted). The reduction operators on a graph
 *        in a node-to-node fashion (n2n). For op == kMean, this means that the aggregation
 *        corresponds to a weighted average.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x (dim_node + dim_node)`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[out] degree          the in-degree of each node, i.e. the sum of incoming
                               edge weights. [on device]
                               [dim = 'graph.n_dst_nodes x 1`]
                               Only required for AggOpT::kMean, may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  edge_weight     the input edge weights [on device]
 *                             [dim = `graph.n_edges x 1`]
 *                             Note that those should we indexed similarly as indices,
 *                             i.e. the edge corresponding to `indices[off]` should have
 *                             the weight `edge_weight[off]`.
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  graph           the input CSR full graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_weighted_n2n_fwd(float* out,
                                 int32_t* out_pos,
                                 float* degree,
                                 const float* node_feat,
                                 const float* edge_weight,
                                 size_t dim_node,
                                 const csc_s32_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);
void agg_concat_weighted_n2n_fwd(float* out,
                                 int64_t* out_pos,
                                 float* degree,
                                 const float* node_feat,
                                 const float* edge_weight,
                                 size_t dim_node,
                                 const csc_s64_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for full-graph aggregation using node features
 *        while concatenting with the orignal node features, the actual aggregation
 *        is done weighting each neighbor according to its edge weight for each
 *        output node (agg_concat_weighted). The reduction operators on a graph
 *        in a node-to-node fashion (n2n). For op == kMean, this means that the aggregation
 *        corresponds to a weighted average.
 *
 * @note `concat_feat` must not overlap with `node_feat` except if
 * `concat_feat == node_feat` and `dim_concat == dim_node`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x (dim_node + dim_node)`]
 * @param[out] out_pos         the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[out] degree          the in-degree of each node, i.e. the sum of incoming
                               edge weights. [on device]
                               [dim = 'graph.n_dst_nodes x 1`]
                               Only required for AggOpT::kMean, may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  concat_feat     the node features to concatenate after aggregation.
 *                             [on device] [dim = `graph.n_dst_nodes x dim_concat`]
 * @param[in]  edge_weight     the input edge weights [on device]
 *                             [dim = `graph.n_edges x 1`]
 *                             Note that those should we indexed similarly as indices,
 *                             i.e. the edge corresponding to `indices[off]` should have
 *                             the weight `edge_weight[off]`.
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  graph           the input CSR full graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_weighted_n2n_fwd(float* out,
                                 int32_t* out_pos,
                                 float* degree,
                                 const float* node_feat,
                                 const float* concat_feat,
                                 const float* edge_weight,
                                 size_t dim_node,
                                 size_t dim_concat,
                                 const bipartite_csc_s32_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);
void agg_concat_weighted_n2n_fwd(float* out,
                                 int64_t* out_pos,
                                 float* degree,
                                 const float* node_feat,
                                 const float* concat_feat,
                                 const float* edge_weight,
                                 size_t dim_node,
                                 size_t dim_concat,
                                 const bipartite_csc_s64_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for full-graph aggregation using node features
 *        while concatenting with the orignal node features, the actual aggregation
 *        is done weighting each neighbor according to its edge weight for each
 *        output node (agg_concat_weighted). The reduction operates on a graph
 *        in a node-to-node fashion (n2n).
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `graph.n_dst_nodes x dim_node`]
 * @param[out] d_edge_weight   the output gradients of edge weights
                               [on device] [dim = `graph.n_edges x 1`]
                               If no gradient required, may be `nullptr`.
 * @param[in] dout             the input node feature gradients. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in] out              the input node feature gradients. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
                               Only required for AggOpT::kMean, may be `nullptr` otherwise.
 * @param[in] out_pos          the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  edge_weight     the input edge weights [on device]
 *                             [dim = `graph.n_edges x 1`]
 *                             Note that those should we indexed similarly as indices,
 *                             i.e. the edge corresponding to `indices[off]` should have
 *                             the weight `edge_weight[off]`.
 * @param[in] degree           the in-degree of each node, i.e. the sum of incoming
                               edge weights. [on device] [dim = 'graph.n_dst_nodes x 1`]
                               Only required for AggOpT::kMean, may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  graph           the input CSR full graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_weighted_n2n_bwd(float* d_node_feat,
                                 float* d_edge_weight,
                                 const float* dout,
                                 const float* out,
                                 const int32_t* out_pos,
                                 const float* edge_weight,
                                 const float* degree,
                                 const float* node_feat,
                                 size_t dim_node,
                                 const csc_s32_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);
void agg_concat_weighted_n2n_bwd(float* d_node_feat,
                                 float* d_edge_weight,
                                 const float* dout,
                                 const float* out,
                                 const int64_t* out_pos,
                                 const float* edge_weight,
                                 const float* degree,
                                 const float* node_feat,
                                 size_t dim_node,
                                 const csc_s64_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);

/**
 * @brief Computes the backward pass for full-graph aggregation using node features
 *        while concatenting with the orignal node features, the actual aggregation
 *        is done weighting each neighbor according to its edge weight for each
 *        output node (agg_concat_weighted). The reduction operates on a graph
 *        in a node-to-node fashion (n2n).
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `graph.n_dst_nodes x dim_node`]
 * @param[out] d_concat_feat   the output gradients of concatenated node features
 *                             [on device] [dim = `graph.n_dst_nodes x dim_concat`].
 *                             Pass `nullptr` if not required.
 * @param[out] d_edge_weight   the output gradients of edge weights
                               [on device] [dim = `graph.n_edges x 1`]
                               If no gradient required, may be `nullptr`.
 * @param[in] dout             the input node feature gradients. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in] out              the input node feature gradients. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
                               Only required for AggOpT::kMean, may be `nullptr` otherwise.
 * @param[in] out_pos          the position of min/max values [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             Only required for AggOpT::kMin and AggOpT::kMax,
 *                             may be `nullptr` otherwise.
 * @param[in]  edge_weight     the input edge weights [on device]
 *                             [dim = `graph.n_edges x 1`]
 *                             Note that those should we indexed similarly as indices,
 *                             i.e. the edge corresponding to `indices[off]` should have
 *                             the weight `edge_weight[off]`.
 * @param[in] degree           the in-degree of each node, i.e. the sum of incoming
                               edge weights. [on device] [dim = 'graph.n_dst_nodes x 1`]
                               Only required for AggOpT::kMean, may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  dim_node        dimensionality of the node features
 * @param[in]  dim_concat      dimensionality of the additional concatenated features
 * @param[in]  graph           the input CSR full graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_weighted_n2n_bwd(float* d_node_feat,
                                 float* d_concat_feat,
                                 float* d_edge_weight,
                                 const float* dout,
                                 const float* out,
                                 const int32_t* out_pos,
                                 const float* edge_weight,
                                 const float* degree,
                                 const float* node_feat,
                                 size_t dim_node,
                                 size_t dim_concat,
                                 const bipartite_csc_s32_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);
void agg_concat_weighted_n2n_bwd(float* d_node_feat,
                                 float* d_concat_feat,
                                 float* d_edge_weight,
                                 const float* dout,
                                 const float* out,
                                 const int64_t* out_pos,
                                 const float* edge_weight,
                                 const float* degree,
                                 const float* node_feat,
                                 size_t dim_node,
                                 size_t dim_concat,
                                 const bipartite_csc_s64_t& graph,
                                 AggOpT op,
                                 const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
