/*
 * 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 multi-head attention layer (GAT-like)
 *        without using cudnn (mha_gat)
 *        operating on full graphs (fg) in a node-to-node
 *        reduction (n2n)
 *
 * @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`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `fg.n_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[out] sm_scores       the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  node_feat       the embeddings of all input nodes. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat [on device]
 *                             [dim = `dim_node + dim_node`]
 * @param[in]  fg              the input CSR graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_s32_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_s64_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_batch_s32_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_fwd(float* out,
                        float* sm_scores,
                        const float* node_feat,
                        const float* attn_weights,
                        const fg_csr_batch_s64_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a multi-head attention layer (GAT-like)
 *        without using cudnn (mha_gat)
 *        operating on full graphs (fg) in a node-to-node
 *        reduction (n2n) but using edge features, too, for
 *        computing the dot product (efeat).
 *
 * @note The `eidx` version uses an edge index to obtain the index of an edge
 * feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` may
 * arbitrary.
 * @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`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `fg.n_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[out] sm_scores       the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  node_feat       the embeddings of all input nodes. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  edge_feat       the embeddings of all edges [on device]
 *                             [dim = `n_edges x dim_edge`]
 *                             For each edge, the features associated with the different
 *                             heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat and edge features [on device]
 *                             [dim = `dim_node + dim_node + dim_edge`]
 * @param[in]  fg              the input CSR graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge        dimension of the edge embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_s32_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_s64_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_batch_s32_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_fwd(float* out,
                              float* sm_scores,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const fg_csr_batch_s64_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_fwd(float* out,
                                   float* sm_scores,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const fg_csr_s32_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_fwd(float* out,
                                   float* sm_scores,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const fg_csr_s64_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_fwd(float* out,
                                   float* sm_scores,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const fg_csr_batch_s32_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_fwd(float* out,
                                   float* sm_scores,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const fg_csr_batch_s64_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a multi-head attention layer (GAT-like)
 *        without using cudnn (mha_gat)
 *        operating on full graphs (fg) in a node-to-node
 *        reduction (n2n)
 *
 * @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`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] grad_node_feat  the gradients on input embeddings of all nodes
 *                             [on device] [dim = `fg.n_nodes x dim_node`].
 * @param[out] grad_weights    the gradients on all attention weights [on device]
 *                             [dim = `dim_node + dim_node`]
 * @param[out] grad_sm_scores  the gradient on softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_out        the gradient on the output of the forward pass
 *                             [on device] [dim = `fg.n_nodes x dim_out`]
 *                             For each node, the gradients associated with
 *                             the different heads are assumed to be
 *                             separate if params.concat_heads = true (`dim_out = dim_node`),
 *                             and shared otherwise (`dim_out = dim_head`).
 * @param[in]  node_feat       the embeddings of all input nodes. [on device]
 *                             [dim = `fg.n_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat [on device]
 *                             [dim = `dim_node + dim_node`]
 * @param[in]  in_sm_scores    the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  fg              the input CSR graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_s32_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_s64_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_batch_s32_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
void mha_gat_fg_n2n_bwd(float* grad_node_feat,
                        float* grad_weights,
                        float* grad_sm_scores,
                        const float* grad_out,
                        const float* node_feat,
                        const float* attn_weights,
                        const float* in_sm_scores,
                        const fg_csr_batch_s64_t& fg,
                        const mha_params& params,
                        int dim_node,
                        const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a multi-head attention layer (GAT-like)
 *        without using cudnn (mha_gat)
 *        operating on full graphs (fg) in a node-to-node
 *        reduction (n2n) but using edge features, too, for
 *        computing the dot product (efeat).
 *
 * @note The `eidx` version uses an edge index to obtain the index of an edge
 * feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` may
 * arbitrary.
 * @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`.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] grad_node_feat  the gradients on input embeddings of all nodes
 *                             [on device] [dim = `fg.n_nodes x dim_node`].
 * @param[out] grad_edge_feat  the gradient on embeddings of all edges [on device]
 *                             [dim = `n_edges x dim_edge`]
 * @param[out] grad_weights    the gradients on all attention weights [on device]
 *                             [dim = `dim_node + dim_node + dim_edge`]
 * @param[in]  grad_sm_scores  the gradient on softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_out        the gradient on the output of the forward pass
 *                             [on device] [dim = `fg.n_nodes x dim_node`]
 * @param[in]  node_feat       the embeddings of all input nodes.
 *                             [on device] [dim = `fg.n_nodes x dim_out`]
 *                             For each node, the gradients associated with
 *                             the different heads are assumed to be
 *                             separate if params.concat_heads = true (`dim_out = dim_node`),
 *                             and shared otherwise (`dim_out = dim_head`).
 * @param[in]  edge_feat       the embeddings of all edges [on device]
 *                             [dim = `n_edges x dim_edge`]
 *                             For each edge, the features associated with the different
 *                             heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat [on device]
 *                             [dim = `dim_node + dim_node + dim_edge`]
 * @param[in]  in_sm_scores    the output softmax scores for backward.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  fg              the input CSR graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge        dimension of the edge embeddings
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_s32_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_s64_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_batch_s32_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_bwd(float* grad_node_feat,
                              float* grad_edge_feat,
                              float* grad_weights,
                              float* grad_sm_scores,
                              const float* grad_out,
                              const float* node_feat,
                              const float* edge_feat,
                              const float* attn_weights,
                              const float* in_sm_scores,
                              const fg_csr_batch_s64_t& fg,
                              const mha_params& params,
                              int dim_node,
                              int dim_edge,
                              const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_bwd(float* grad_node_feat,
                                   float* grad_edge_feat,
                                   float* grad_weights,
                                   float* grad_sm_scores,
                                   const float* grad_out,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const float* in_sm_scores,
                                   const fg_csr_s32_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_bwd(float* grad_node_feat,
                                   float* grad_edge_feat,
                                   float* grad_weights,
                                   float* grad_sm_scores,
                                   const float* grad_out,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const float* in_sm_scores,
                                   const fg_csr_s64_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_bwd(float* grad_node_feat,
                                   float* grad_edge_feat,
                                   float* grad_weights,
                                   float* grad_sm_scores,
                                   const float* grad_out,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const float* in_sm_scores,
                                   const fg_csr_batch_s32_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
void mha_gat_fg_n2n_efeat_eidx_bwd(float* grad_node_feat,
                                   float* grad_edge_feat,
                                   float* grad_weights,
                                   float* grad_sm_scores,
                                   const float* grad_out,
                                   const float* node_feat,
                                   const float* edge_feat,
                                   const float* attn_weights,
                                   const float* in_sm_scores,
                                   const fg_csr_batch_s64_t& fg,
                                   const mha_params& params,
                                   int dim_node,
                                   int dim_edge,
                                   const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
