/*
 * 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/aggregator.hpp>
#include <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/graph/format.hpp>

#include <cstdint>

namespace cugraph::ops::gnn {

/**
 * @brief Forward pass for fused aggregation and concat (self features) layer
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `mfg.n_out_nodes x 2dim`]
 * @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.
 *                             It must be initialized.
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_fwd(float* out,
                    int32_t* out_pos,
                    const float* in,
                    size_t dim,
                    const mfg_ellpack_s32_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
void agg_concat_fwd(float* out,
                    int64_t* out_pos,
                    const float* in,
                    size_t dim,
                    const mfg_ellpack_s64_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
void agg_concat_fwd(float* out,
                    int32_t* out_pos,
                    const float* in,
                    size_t dim,
                    const mfg_csr_s32_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
void agg_concat_fwd(float* out,
                    int64_t* out_pos,
                    const float* in,
                    size_t dim,
                    const mfg_csr_s64_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
/** @} */

/**
 * @brief Backward pass for fused aggregation and concat (self features) layer
 *
 * @param[out] 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 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             dimensionality of the embeddings
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 *                             It must be initialized.
 * @param[in]  op              aggregation operation
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_concat_bwd(float* din,
                    const float* dout,
                    const int32_t* out_pos,
                    size_t dim,
                    const mfg_ellpack_s32_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
void agg_concat_bwd(float* din,
                    const float* dout,
                    const int64_t* out_pos,
                    size_t dim,
                    const mfg_ellpack_s64_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
void agg_concat_bwd(float* din,
                    const float* dout,
                    const int32_t* out_pos,
                    size_t dim,
                    const mfg_csr_s32_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
void agg_concat_bwd(float* din,
                    const float* dout,
                    const int64_t* out_pos,
                    size_t dim,
                    const mfg_csr_s64_t& mfg,
                    AggOpT op,
                    const cuda::stream& stream);
/** @} */

/**
 * @brief Backward pass for fused aggregation and concat (self features) layer
 *        using reverse graph
 *
 * 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[out] 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 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             dimensionality of the embeddings
 * @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]  stream          cuda stream
 * @{
 */
void agg_concat_bwd_rev(float* din,
                        const float* dout,
                        const int32_t* out_pos,
                        size_t dim,
                        const mfg_ellpack_s32_t& mfg,
                        const mfg_csr_rev_s32_t& mfg_rev,
                        AggOpT op,
                        const cuda::stream& stream);
void agg_concat_bwd_rev(float* din,
                        const float* dout,
                        const int64_t* out_pos,
                        size_t dim,
                        const mfg_ellpack_s64_t& mfg,
                        const mfg_csr_rev_s64_t& mfg_rev,
                        AggOpT op,
                        const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
