/*
 * 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 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 full graph (fg)
 *        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 = `fg.n_nodes x (dim_node + 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[out] degree          the in-degree of each node, i.e. the sum of incoming
                               edge weights. [on device]
                               [dim = 'fg.n_nodes x 1`]
                               Only required for AggOpT::kMean, may be `nullptr` otherwise.
 * @param[in]  node_feat       the input node features. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 * @param[in]  edge_weight     the input edge weights [on device]
 *                             [dim = `fg.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]  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_concat_weighted_fg_n2n_fwd(float* out,
                                    int32_t* out_pos,
                                    float* degree,
                                    const float* node_feat,
                                    const float* edge_weight,
                                    size_t dim_node,
                                    const fg_csr_s32_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_concat_weighted_fg_n2n_fwd(float* out,
                                    int32_t* out_pos,
                                    float* degree,
                                    const float* node_feat,
                                    const float* edge_weight,
                                    size_t dim_node,
                                    const fg_csr_batch_s32_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_concat_weighted_fg_n2n_fwd(float* out,
                                    int64_t* out_pos,
                                    float* degree,
                                    const float* node_feat,
                                    const float* edge_weight,
                                    size_t dim_node,
                                    const fg_csr_s64_t& fg,
                                    AggOpT op,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_concat_weighted_fg_n2n_fwd(float* out,
                                    int64_t* out_pos,
                                    float* degree,
                                    const float* node_feat,
                                    const float* edge_weight,
                                    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 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 full graph (fg)
 *        in a node-to-node fashion (n2n).
 *
 * @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_weight   the output gradients of edge weights
                               [on device] [dim = `fg.n_edges x 1`]
                               Only required if `need_ew_grad`, may be nullptr otherwise.
 * @param[in] dout             the input node feature gradients. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 * @param[in] out              the input node feature gradients. [on device]
 *                             [dim = `fg.n_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 = `fg.n_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 = `fg.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 = 'fg.n_nodes x 1`]
                               Only required for AggOpT::kMean, 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]  need_ew_grad    if gradient on edge weights is needed
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_weighted_fg_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 fg_csr_s32_t& fg,
                                    AggOpT op,
                                    bool need_ew_grad,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_concat_weighted_fg_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 fg_csr_batch_s32_t& fg,
                                    AggOpT op,
                                    bool need_ew_grad,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_concat_weighted_fg_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 fg_csr_s64_t& fg,
                                    AggOpT op,
                                    bool need_ew_grad,
                                    bool ignore_self,
                                    const cuda::stream& stream);
void agg_concat_weighted_fg_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 fg_csr_batch_s64_t& fg,
                                    AggOpT op,
                                    bool need_ew_grad,
                                    bool ignore_self,
                                    const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
