/*
 * 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 D-MPNN-like full graph aggregation
 *        (agg_dmpnn_fg) using edge features in an edge-to-edge reduction
 *        scheme (e2e)
 *
 * @note For these operations, if the graph uses edge indices
 *       (`fg.ef_indices != nullptr`), then these indices must be a partition /
 *       permutation of the edge IDs.
 *
 * @param[out] out_edge_feat   the output edge features [on device]
 *                             output dimension `out_dim` is `dim_edge * 2` if
 *                             `concat_own`, otherwise `dim_edge`.
 * @param[in]  in_edge_feat    the input edge features. [on device]
 *                             [dim = `fg.n_indices x dim_edge`].
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  fg              the input CSR full graph. It must be undirected:
 *                             for every edge, there is a reverse edge present
 *                             in the structure, too.
 *                             This can be a batch of (undirected) graphs, too.
 * @param[in]  concat_own      If `true`, concatenate an edge's input features
 *                             to the aggregated output features.
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_dmpnn_fg_e2e_fwd(float* out_edge_feat,
                          const float* in_edge_feat,
                          size_t dim_edge,
                          const fg_csr_s32_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_fwd(float* out_edge_feat,
                          const float* in_edge_feat,
                          size_t dim_edge,
                          const fg_csr_batch_s32_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_fwd(float* out_edge_feat,
                          const float* in_edge_feat,
                          size_t dim_edge,
                          const fg_csr_s64_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_fwd(float* out_edge_feat,
                          const float* in_edge_feat,
                          size_t dim_edge,
                          const fg_csr_batch_s64_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_fwd(float* out_edge_feat,
                               const float* in_edge_feat,
                               size_t dim_edge,
                               const fg_csr_s32_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_fwd(float* out_edge_feat,
                               const float* in_edge_feat,
                               size_t dim_edge,
                               const fg_csr_batch_s32_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_fwd(float* out_edge_feat,
                               const float* in_edge_feat,
                               size_t dim_edge,
                               const fg_csr_s64_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_fwd(float* out_edge_feat,
                               const float* in_edge_feat,
                               size_t dim_edge,
                               const fg_csr_batch_s64_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for D-MPNN-like full graph aggregation
 *        (agg_dmpnn_fg) using edge features in an edge-to-edge reduction
 *        scheme (e2e)
 *
 * @note For these operations, if the graph uses edge indices
 *       (`fg.ef_indices != nullptr`), then these indices must be a partition /
 *       permutation of the edge IDs.
 *
 * @param[out] d_in_edge_feat  the output gradients on input edge features
 *                             [on device] [dim = `fg.n_indices x dim_edge`].
 * @param[in]  d_out_edge_feat the input gradients on output edge features.
 *                             [on device] [dim = `fg.n_indices x out_dim`].
 *                             For `out_dim`, see parameter `out_edge_feat`
 *                             of the forward pass.
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  fg              the input CSR full graph. It must be undirected:
 *                             for every edge, there is a reverse edge present
 *                             in the structure, too.
 *                             This can be a batch of (undirected) graphs, too.
 * @param[in]  concat_own      If `true`, concatenate an edge's input features
 *                             to the aggregated output features.
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_dmpnn_fg_e2e_bwd(float* d_in_edge_feat,
                          const float* d_out_edge_feat,
                          size_t dim_edge,
                          const fg_csr_s32_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_bwd(float* d_in_edge_feat,
                          const float* d_out_edge_feat,
                          size_t dim_edge,
                          const fg_csr_batch_s32_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_bwd(float* d_in_edge_feat,
                          const float* d_out_edge_feat,
                          size_t dim_edge,
                          const fg_csr_s64_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_bwd(float* d_in_edge_feat,
                          const float* d_out_edge_feat,
                          size_t dim_edge,
                          const fg_csr_batch_s64_t& fg,
                          bool concat_own,
                          const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_bwd(float* d_in_edge_feat,
                               const float* d_out_edge_feat,
                               size_t dim_edge,
                               const fg_csr_s32_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_bwd(float* d_in_edge_feat,
                               const float* d_out_edge_feat,
                               size_t dim_edge,
                               const fg_csr_batch_s32_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_bwd(float* d_in_edge_feat,
                               const float* d_out_edge_feat,
                               size_t dim_edge,
                               const fg_csr_s64_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
void agg_dmpnn_fg_e2e_eidx_bwd(float* d_in_edge_feat,
                               const float* d_out_edge_feat,
                               size_t dim_edge,
                               const fg_csr_batch_s64_t& fg,
                               bool concat_own,
                               const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
