/*
 * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#pragma once

#include <cugraph-ops/activation.hpp>
#include <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/operators/common.hpp>
#include <cugraph-ops/test/graph/format_mfg.hpp>

#include <memory>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n).
 *
 * @note We always assume `dim_head = dim / params.num_heads`.
 * @note If `params.concat_heads = true`, we have `dim_out = dim`, otherwise
 * `dim_out = dim_head`
 *
 * @param[out] attn_feat       the output embeddings after attention.
 *                             [on device] [dim = `mfg.n_out_nodes x dim_out`].
 *                             If `params.concat_heads = true`, the features
 *                             associated with the different heads are
 *                             concatenated, with total size being `dim`,
 *                             otherwise they are averaged with total size being
 *                             `dim_head`.
 * @param[inout] sm_scores     the output softmax scores for the backward pass.
 *                             [on device] [dim = `2 x mfg.n_out_nodes x
 *                             params.num_heads x mfg.sample_size`]
 * @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.
 * @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& graph,
                         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& graph,
                         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& graph,
                         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& graph,
                         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_csc_s32_t& graph,
                         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_csc_s64_t& graph,
                         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_csc_s32_t& graph,
                         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_csc_s64_t& graph,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass of a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on message flow graphs (mfg) in a node-to-node
 *        reduction (n2n).
 *
 * @note We always assume `dim_head = dim / params.num_heads`.
 * @note If `params.concat_heads = true`, we have `dim_out = dim`, otherwise
 * `dim_out = dim_head`
 *
 * @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_out`]
 * @param[in]  in_feat         the embeddings of all input nodes from the forward pass.
 *                             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[inout]  sm_scores    the softmax scores from the forward pass.
 *                             [on device] [dim = `2 x mfg.n_out_nodes x
 *                             params.num_heads x mfg.sample_size`]
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 * @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& graph,
                         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& graph,
                         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& graph,
                         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& graph,
                         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_csc_s32_t& graph,
                         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_csc_s64_t& graph,
                         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_csc_s32_t& graph,
                         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_csc_s64_t& graph,
                         const mha_params& params,
                         size_t dim,
                         const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
