/*
 * Copyright (c) 2021-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/activation.hpp>
#include <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/graph/format.hpp>
#include <cugraph-ops/operators/common.hpp>

#include <memory>

namespace cugraph::ops::gnn {

/**
 * @brief Computes the forward pass for a multi-head attention layer (GAT-like)
 *        without using cudnn (mha_gat)
 *        operating on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n)
 *
 * @param[out] attn_feat       the output embeddings after attention.
 *                             [on device] [dim = `mfg.n_out_nodes x dim`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated, with total
 *                             size being `dim`.
 * @param[out] sm_scores       the output softmax scores for backward.
 *                             [on device] [dim = `2 x mfg.n_out_nodes x
 *                             params.num_heads x (mfg.sample_size + 1)`]
 * @param[in]  in_feat         the embeddings of all input nodes. This must
 *                             include embeddings of output / destination nodes
 *                             [on device] [dim = `*mfg.n_in_nodes x dim`].
 *                             For each node, the features associated with
 *                             the different heads are expected to be
 *                             concatenated, with total size being `dim`.
 * @param[in]  weights         the weights applied to in/out feat [on device]
 *                             [dim = `dim + dim`]
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG. It must be initialized.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim             dimension of the input and output embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_mfg_n2n_fwd(float* attn_feat,
                         float* sm_scores,
                         const float* in_feat,
                         const float* weights,
                         const mfg_ellpack_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_fwd(float* attn_feat,
                         float* sm_scores,
                         const float* in_feat,
                         const float* weights,
                         const mfg_ellpack_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_fwd(double* attn_feat,
                         double* sm_scores,
                         const double* in_feat,
                         const double* weights,
                         const mfg_ellpack_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_fwd(double* attn_feat,
                         double* sm_scores,
                         const double* in_feat,
                         const double* weights,
                         const mfg_ellpack_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_fwd(float* attn_feat,
                         float* sm_scores,
                         const float* in_feat,
                         const float* weights,
                         const mfg_csr_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_fwd(float* attn_feat,
                         float* sm_scores,
                         const float* in_feat,
                         const float* weights,
                         const mfg_csr_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_fwd(double* attn_feat,
                         double* sm_scores,
                         const double* in_feat,
                         const double* weights,
                         const mfg_csr_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_fwd(double* attn_feat,
                         double* sm_scores,
                         const double* in_feat,
                         const double* weights,
                         const mfg_csr_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass of a multi-head attention layer (GAT-like)
 *        without using cudnn (mha_gat)
 *        operating on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n)
 *
 * @param[out] d_in_feat       the data gradient on the input features
 *                             (must be a separate buffer than in_feat).
 *                             [on device] [dim = `*mfg.n_in_nodes x dim`]
 * @param[out] d_weights       the weight gradient on the weights.
 *                             (must be a separate buffer than weights).
 *                             [on device] [dim = `dim + dim`]
 * @param[in]  d_attn_feat     the input data gradients of the out node features
 *                             [on device] [dim = `mfg.n_out_nodes x dim`]
 * @param[in]  in_feat         the embeddings of all input nodes from forward.
 *                             This must include embeddings of output /
 *                             destination nodes
 *                             [on device] [dim = `*mfg.n_in_nodes x dim`]
 * @param[in]  weights         the weights applied to in/out features
 *                             [on device] [dim = `dim + dim`]
 * @param[in]  sm_scores       the softmax scores from forward.
 *                             [on device] [dim = `2 x mfg.n_out_nodes x
 *                             params.num_heads x (mfg.sample_size + 1)`]
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG. It must be initialized.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim             dimension of the input and output embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_mfg_n2n_bwd(float* d_in_feat,
                         float* d_weights,
                         const float* d_attn_feat,
                         const float* in_feat,
                         const float* weights,
                         const float* sm_scores,
                         const mfg_ellpack_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_bwd(float* d_in_feat,
                         float* d_weights,
                         const float* d_attn_feat,
                         const float* in_feat,
                         const float* weights,
                         const float* sm_scores,
                         const mfg_ellpack_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_bwd(double* d_in_feat,
                         double* d_weights,
                         const double* d_attn_feat,
                         const double* in_feat,
                         const double* weights,
                         const double* sm_scores,
                         const mfg_ellpack_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_bwd(double* d_in_feat,
                         double* d_weights,
                         const double* d_attn_feat,
                         const double* in_feat,
                         const double* weights,
                         const double* sm_scores,
                         const mfg_ellpack_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_bwd(float* d_in_feat,
                         float* d_weights,
                         const float* d_attn_feat,
                         const float* in_feat,
                         const float* weights,
                         const float* sm_scores,
                         const mfg_csr_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_bwd(float* d_in_feat,
                         float* d_weights,
                         const float* d_attn_feat,
                         const float* in_feat,
                         const float* weights,
                         const float* sm_scores,
                         const mfg_csr_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_bwd(double* d_in_feat,
                         double* d_weights,
                         const double* d_attn_feat,
                         const double* in_feat,
                         const double* weights,
                         const double* sm_scores,
                         const mfg_csr_s32_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
void mha_gat_mfg_n2n_bwd(double* d_in_feat,
                         double* d_weights,
                         const double* d_attn_feat,
                         const double* in_feat,
                         const double* weights,
                         const double* sm_scores,
                         const mfg_csr_s64_t& mfg,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
