/*
 * 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 a weighted aggregation (agg_weighted)
 *        using node features in an node-to-node reduction (n2n)
 *        operating on a graph where each neighbor is weighted according
 *        to its edge weight. For op == kMean, the mean corresponds to the weighted
 *        average per dimension given the set of neighbors.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x 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 graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_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_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);

void agg_weighted_n2n_fwd(float* out,
                          int32_t* out_pos,
                          float* degree,
                          const float* node_feat,
                          const float* edge_weight,
                          size_t dim_node,
                          const bipartite_csc_s32_t& graph,
                          AggOpT op,
                          const cuda::stream& stream);
void agg_weighted_n2n_fwd(float* out,
                          int64_t* out_pos,
                          float* degree,
                          const float* node_feat,
                          const float* edge_weight,
                          size_t dim_node,
                          const bipartite_csc_s64_t& graph,
                          AggOpT op,
                          const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a weighted aggregation (agg_weighted)
 *        operating on a graph in a node-to-node reduction (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 gradient not required, may be `nullptr`.
 * @param[in]  dout            the input node feature gradients. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 * @param[out] 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[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[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[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]  dim_node        dimensionality of the node features
 * @param[in]  graph           the input graph
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_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_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);

void agg_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 bipartite_csc_s32_t& graph,
                          AggOpT op,
                          const cuda::stream& stream);
void agg_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 bipartite_csc_s64_t& graph,
                          AggOpT op,
                          const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
