/*
 * 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 a vanilla aggregation (agg_simple)
 *        operatoring on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n)
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `mfg.n_out_nodes x dim`]
 * @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`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `*mfg.n_in_nodes x dim`]
 * @param[in]  dim             dimensionality of the embeddings
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_mfg_n2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* in,
                            size_t dim,
                            const mfg_ellpack_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* in,
                            size_t dim,
                            const mfg_ellpack_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* in,
                            size_t dim,
                            const mfg_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* in,
                            size_t dim,
                            const mfg_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a vanilla aggregation (agg_simple)
 *        operatoring on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n)
 *
 * @param[out] din             output gradient of the embeddings of this layer.
 *                             [on device] [dim = `*mfg.n_in_nodes x dim`]
 * @param[in]  dout            the input gradient of the embeddings from the
 *                             layer above. [on device]
 *                             [dim = `mfg.n_out_nodes x dim`]
 * @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             dimensionality of the embeddings
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim,
                            const mfg_ellpack_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim,
                            const mfg_ellpack_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim,
                            const mfg_csc_s32_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim,
                            const mfg_csc_s64_t& graph,
                            AggOpT op,
                            const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a vanilla aggregation (agg_simple)
 *        operatoring 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.
 *
 * @param[in]  din             output gradient of the embeddings of this layer.
 *                             [on device] [dim = `*mfg.n_in_nodes x dim`]
 * @param[out] dout            the input gradient of the embeddings from the
 *                             layer above. [on device]
 *                             [dim = `mfg.n_out_nodes x dim`]
 * @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             dimensionality of the 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_simple_mfg_n2n_bwd_rev(float* din,
                                const float* dout,
                                const int32_t* out_pos,
                                size_t dim,
                                const mfg_ellpack_s32_t& graph,
                                const mfg_csc_rev_s32_t& mfg_rev,
                                AggOpT op,
                                const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd_rev(float* din,
                                const float* dout,
                                const int64_t* out_pos,
                                size_t dim,
                                const mfg_ellpack_s64_t& graph,
                                const mfg_csc_rev_s64_t& mfg_rev,
                                AggOpT op,
                                const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
