/*
 * 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/operators/common.hpp>
#include <cugraph-ops/test/graph/format_mfg.hpp>

#include <cstdint>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for fused aggregation and concat (additional features)
 *        layer (agg_concat) operating on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n)
 *
 * @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 = `mfg.n_out_nodes x (dim_node + dim_concat)`]
 * @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 = `mfg.n_out_nodes x dim_node`]
 * @param[in]  node_feat       the input node embeddings. [on device]
 *                             [dim = `mfg.n_in_nodes x dim_node`]
 * @param[in]  concat_feat     the input additional embeddings. [on device]
 *                             [dim = `mfg.n_in_nodes x dim_concat`]
 * @param[in]  dim_node        dimensionality of the node embeddings
 * @param[in]  dim_concat      dimensionality of the additional embeddings
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_mfg_n2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* node_feat,
                            const float* concat_feat,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_ellpack_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_concat_mfg_n2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* node_feat,
                            const float* concat_feat,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_ellpack_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_concat_mfg_n2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* node_feat,
                            const float* concat_feat,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_concat_mfg_n2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* node_feat,
                            const float* concat_feat,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for fused aggregation and concat (additional features)
 *        layer (agg_concat) operating on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n)
 *
 * @note Dimensions passed must be the same as in forward, regardless of whether
 *       the gradient for a given feature is actually required or not.
 * @note `d_concat_feat` must not overlap with `d_node_feat` except if
 * `d_concat_feat == d_node_feat` and `dim_concat == dim_node`.
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `mfg.n_in_nodes x dim_node`].
 *                             Pass `nullptr` if not required.
 * @param[out] d_concat_feat   the output gradients of concatenated node features
 *                             [on device] [dim = `mfg.n_in_nodes x dim_concat`].
 *                             Pass `nullptr` if not required.
 * @param[out] dout            the input gradient of the embeddings from the
 *                             layer above. [on device]
 *                             [dim = `mfg.n_out_nodes x 2dim`]
 * @param[in]  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 = `mfg.n_out_nodes x dim`]
 * @param[in]  dim_node        dimensionality of the node embeddings
 * @param[in]  dim_concat      dimensionality of the additional embeddings
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_mfg_n2n_bwd(float* d_node_feat,
                            float* d_concat_feat,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_ellpack_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_concat_mfg_n2n_bwd(float* d_node_feat,
                            float* d_concat_feat,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_ellpack_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_concat_mfg_n2n_bwd(float* d_node_feat,
                            float* d_concat_feat,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_concat_mfg_n2n_bwd(float* d_node_feat,
                            float* d_concat_feat,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim_node,
                            size_t dim_concat,
                            const mfg_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for fused aggregation and concat (additional features)
 *        layer (agg_concat) operating on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n) using the reverse graph (bwd_rev)
 *
 * This avoids using atomics, but is unusable performance-wise due to the large
 * overheads of sorting (as of now), which is necessary for correctness
 * in the case of multi-graphs and min/max aggregation.
 *
 * @note Dimensions passed must be the same as in forward, regardless of whether
 *       the gradient for a given feature is actually required or not.
 * @note `d_concat_feat` must not overlap with `d_node_feat` except if
 * `d_concat_feat == d_node_feat` and `dim_concat == dim_node`.
 *
 * @param[out] d_node_feat     the output gradients of input node features
 *                             [on device] [dim = `mfg.n_in_nodes x dim_node`].
 *                             Pass `nullptr` if not required.
 * @param[out] d_concat_feat   the output gradients of concatenated node features
 *                             [on device] [dim = `mfg.n_in_nodes x dim_concat`].
 *                             Pass `nullptr` if not required.
 * @param[out] dout            the input gradient of the embeddings from the
 *                             layer above. [on device]
 *                             [dim = `mfg.n_out_nodes x 2dim`]
 * @param[in]  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 = `mfg.n_out_nodes x dim`]
 * @param[in]  dim_node        dimensionality of the node embeddings
 * @param[in]  dim_concat      dimensionality of the additional embeddings
 * @param[in]  mfg             the input (ELLPACK) MFG.
 * @param[in]  mfg_rev         the input reverse (CSR) MFG.
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_mfg_n2n_bwd_rev(float* d_node_feat,
                                float* d_concat_feat,
                                const float* dout,
                                const int32_t* out_pos,
                                size_t dim_node,
                                size_t dim_concat,
                                const mfg_ellpack_s32_t& graph,
                                const mfg_csc_rev_s32_t& mfg_rev,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_concat_mfg_n2n_bwd_rev(float* d_node_feat,
                                float* d_concat_feat,
                                const float* dout,
                                const int64_t* out_pos,
                                size_t dim_node,
                                size_t dim_concat,
                                const mfg_ellpack_s64_t& graph,
                                const mfg_csc_rev_s64_t& mfg_rev,
                                AggOpT op,
                                const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
