/*
 * 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 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 ld`]
 * @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]  ld              leading dimension for `out` array. It must be
 *                             at least `dim`.
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 *                             It must be initialized.
 * @param[in]  op              aggregation operation
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_mfg_n2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* in,
                            size_t dim,
                            size_t ld,
                            const mfg_ellpack_s32_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* in,
                            size_t dim,
                            size_t ld,
                            const mfg_ellpack_s64_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_fwd(float* out,
                            int32_t* out_pos,
                            const float* in,
                            size_t dim,
                            size_t ld,
                            const mfg_csr_s32_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_fwd(float* out,
                            int64_t* out_pos,
                            const float* in,
                            size_t dim,
                            size_t ld,
                            const mfg_csr_s64_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            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 ld`]
 * @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]  ld              leading dimension for `dout` and `out` arrays.
 *                             It must be at least `dim`.
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 *                             It must be initialized.
 * @param[in]  op              aggregation operation
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim,
                            size_t ld,
                            const mfg_ellpack_s32_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim,
                            size_t ld,
                            const mfg_ellpack_s64_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int32_t* out_pos,
                            size_t dim,
                            size_t ld,
                            const mfg_csr_s32_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd(float* din,
                            const float* dout,
                            const int64_t* out_pos,
                            size_t dim,
                            size_t ld,
                            const mfg_csr_s64_t& mfg,
                            AggOpT op,
                            bool ignore_self,
                            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 ld`]
 * @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]  ld              leading dimension for `dout` and `out` arrays.
 *                             It must be at least `dim`.
 * @param[in]  mfg             the input (ELLPACK) MFG. It must be initialized.
 * @param[in]  mfg_rev         the input reverse (CSR) MFG.
 *                             It must be initialized.
 * @param[in]  op              aggregation operation
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @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,
                                size_t ld,
                                const mfg_ellpack_s32_t& mfg,
                                const mfg_csr_rev_s32_t& mfg_rev,
                                AggOpT op,
                                bool ignore_self,
                                const cuda::stream& stream);
void agg_simple_mfg_n2n_bwd_rev(float* din,
                                const float* dout,
                                const int64_t* out_pos,
                                size_t dim,
                                size_t ld,
                                const mfg_ellpack_s64_t& mfg,
                                const mfg_csr_rev_s64_t& mfg_rev,
                                AggOpT op,
                                bool ignore_self,
                                const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
