/*
 * Copyright (c) 2020-2023, 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 concetenation of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_cat) for a full graph (fg)
 * in an edge-to-edge fashion (e2e)
 *
 * @note If any of `dim_edge`, `dim_src`, or `dim_dst` is set to `0`, the corresponding
 * device array is not considered for this operation. In those cases, the corresponding
 * feature pointer, e.g. `edge_feat`, can be passed as `nullptr`.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 *
 * @param[out] out             the output embeddings after concatenation.
 *                             [on device] [dim = n_edges x (dim_edge + dim_src + dim_dst)`].
 * @param[in]  edge_feat       the input edge embeddings [on device]
 *                             [dim = `n_edges x dim_edge`]
 * @param[in]  src_feat        the input source-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_src`]
 * @param[in]  dst_feat        the input destination-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_dst`]
 * @param[in]  fg              the input CSR graph.
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  dim_src         dimension of the input source node embeddings.
 * @param[in]  dim_dst         dimension of the input destinaton node embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_fg_e2e_concat_fwd(float* out,
                                    const float* edge_feat,
                                    const float* src_feat,
                                    const float* dst_feat,
                                    const fg_csr_s32_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
void update_efeat_fg_e2e_concat_fwd(float* out,
                                    const float* edge_feat,
                                    const float* src_feat,
                                    const float* dst_feat,
                                    const fg_csr_batch_s32_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
void update_efeat_fg_e2e_concat_fwd(float* out,
                                    const float* edge_feat,
                                    const float* src_feat,
                                    const float* dst_feat,
                                    const fg_csr_s64_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
void update_efeat_fg_e2e_concat_fwd(float* out,
                                    const float* edge_feat,
                                    const float* src_feat,
                                    const float* dst_feat,
                                    const fg_csr_batch_s64_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a summation of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_sum) for a full graph (fg)
 * in an edge-to-edge fashion (e2e)
 *
 *
 * @note If any pointers `edge_feat`, `src_feat`, or `dst_feat` is passed as `nullptr`,
 * the corresponding features are not considered for the summmation operation.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 *
 * @param[out] out             the output embeddings after summation.
 *                             [on device] [dim = n_edges x dim_edge`].
 * @param[in]  edge_feat       the input edge embeddings [on device]
 *                             [dim = `n_edges x dim_edge`]
 * @param[in]  src_feat        the input source-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_edge`]
 * @param[in]  dst_feat        the input destination-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_edge`]
 * @param[in]  fg              the input CSR graph.
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_fg_e2e_sum_fwd(float* out,
                                 const float* edge_feat,
                                 const float* src_feat,
                                 const float* dst_feat,
                                 const fg_csr_s32_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
void update_efeat_fg_e2e_sum_fwd(float* out,
                                 const float* edge_feat,
                                 const float* src_feat,
                                 const float* dst_feat,
                                 const fg_csr_batch_s32_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
void update_efeat_fg_e2e_sum_fwd(float* out,
                                 const float* edge_feat,
                                 const float* src_feat,
                                 const float* dst_feat,
                                 const fg_csr_s64_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
void update_efeat_fg_e2e_sum_fwd(float* out,
                                 const float* edge_feat,
                                 const float* src_feat,
                                 const float* dst_feat,
                                 const fg_csr_batch_s64_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for the concat-op of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_cat) for a full graph (fg)
 * in an edge-to-edge fashion (e2e)
 *
 * @note If any of `dim_edge`, `dim_src`, or `dim_dst` is set to `0` in the forward pass,
 * the corresponding device array is also not considered in the backward pass. In any case,
 * the dimensions are expected to be the same as in the forward pass.
 * @note If a gradient w.r.t. a specific input embedding, e.g. `grad_edge_feat`,
 * is not needed, the pointer can be passed as `nullptr` and the gradient computation
 * will be skipped. Have in mind that the corresponding dimension still might be necessary,
 * see note above and forward for details on that.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 *
 * @param[out] grad_edge_feat  the output gradient on the edge embeddings [on device]
 *                             [dim = `n_edges x dim_edge`]
 * @param[out] grad_src_feat   the output gradient on the source-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_src`]
 * @param[out] grad_dst_feat   the output gradient on the destionation-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_dst`]
 * @param[in]  grad_out        the input gradient ofn the output embeddings after concatenation.
 *                             [on device] [dim = n_edges x (dim_edge + dim_src + dim_dst)`].
 * @param[in]  fg              the input CSR graph.
 * @param[in]  dim_edge        dimension of the input edge embeddings.
 * @param[in]  dim_src         dimension of the input source node embeddings.
 * @param[in]  dim_dst         dimension of the input destinaton node embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_fg_e2e_concat_bwd(float* grad_edge_feat,
                                    float* grad_src_feat,
                                    float* grad_dst_feat,
                                    const float* grad_out,
                                    const fg_csr_s32_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
void update_efeat_fg_e2e_concat_bwd(float* grad_edge_feat,
                                    float* grad_src_feat,
                                    float* grad_dst_feat,
                                    const float* grad_out,
                                    const fg_csr_batch_s32_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
void update_efeat_fg_e2e_concat_bwd(float* grad_edge_feat,
                                    float* grad_src_feat,
                                    float* grad_dst_feat,
                                    const float* grad_out,
                                    const fg_csr_s64_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
void update_efeat_fg_e2e_concat_bwd(float* grad_edge_feat,
                                    float* grad_src_feat,
                                    float* grad_dst_feat,
                                    const float* grad_out,
                                    const fg_csr_batch_s64_t& fg,
                                    size_t dim_edge,
                                    size_t dim_src,
                                    size_t dim_dst,
                                    const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a sum-op of edge features
 * features of the source nodes, and the features of the target nodes
 * to create new edge features (update_efeat_sum) for a full graph (fg)
 * in an edge-to-edge fashion (e2e)
 *
 *
 * @note If a gradient w.r.t. a specific input embedding, e.g. `grad_edge_feat`,
 * is not needed, the pointer can be passed as `nullptr` and the gradient computation
 * will be skipped.
 * @note `dim_edge` is expected to be the same as in the forward pass.
 * @note When specifying shapes of buffers, we'll refer to
 * both `fg.n_indices` for a single graph or `fg_batch.n_edges`
 * for a batch of graphs simply as `n_edges`.
 *
 * @param[out] grad_edge_feat  the output gradient on the edge embeddings
 *                             [on device][dim = `n_edges x dim_edge`]
 * @param[out] grad_src_feat   the output gradient on the source-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_edge`]
 * @param[out] grad_dst_feat   the output gradient on the destination-node embeddings
 *                             [on device] [dim = `fg.n_nodes x dim_edge`]
 * @param[in]  grad_out        the input gradient on the output embeddings after summation.
 *                             [on device] [dim = n_edges x dim_edge`].
 * @param[in]  fg              the input CSR graph.
 * @param[in]  dim_edge        dimension of the input embeddings.
 * @param[in]  stream          cuda stream
 * @{
 */
void update_efeat_fg_e2e_sum_bwd(float* grad_edge_feat,
                                 float* grad_src_feat,
                                 float* grad_dst_feat,
                                 const float* grad_out,
                                 const fg_csr_s32_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
void update_efeat_fg_e2e_sum_bwd(float* grad_edge_feat,
                                 float* grad_src_feat,
                                 float* grad_dst_feat,
                                 const float* grad_out,
                                 const fg_csr_batch_s32_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
void update_efeat_fg_e2e_sum_bwd(float* grad_edge_feat,
                                 float* grad_src_feat,
                                 float* grad_dst_feat,
                                 const float* grad_out,
                                 const fg_csr_s64_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
void update_efeat_fg_e2e_sum_bwd(float* grad_edge_feat,
                                 float* grad_src_feat,
                                 float* grad_dst_feat,
                                 const float* grad_out,
                                 const fg_csr_batch_s64_t& fg,
                                 size_t dim_edge,
                                 const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
