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

#include <cstdint>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for a heterogenous aggregation (RGCN-like)
 *        potentially using a basis decomposition (agg_hg_basis)
 *        operating on a graph in a node-to-node reduction (n2n)
 *
 * @note This is based on basis decomposition and assumes that output node
 *       features are concatenated by basis matrix features, i.e. the basis
 *       matrix multiplication is performed before aggregation.
 *       In particular, we define `dim` and `out_dim` as follows:
 *       if (weights_comb != nullptr && !concat_own) dim = n_bases x out_dim
 *       if (weights_comb != nullptr && concat_own)  dim = (n_bases + 1) x out_dim
 *       if (weights_comb == nullptr && !concat_own) dim = graph.n_edge_types x out_dim
 *       if (weights_comb == nullptr && concat_own)  dim = (graph.n_edge_types + 1) x out_dim
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x out_dim`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `graph.n_edge_types x n_bases`].
 *                             This may be `nullptr` if the combination will be
 *                             included in basis matrix multiplication.
 *                             In this case, `n_bases` is ignored.
 * @param[in]  dim             dimensionality of the embeddings.
 *                             See above how it is defined w.r.t. `out_dim`.
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  graph           the input graph
 * @param[in]  concat_own      If `true`, we assume that node features were
 *                             concatenated to the input embeddings after all
 *                             basis features. In this case, these features
 *                             are aggregated into the output features
 *                             with a weight of `1`.
 * @param[in]  norm_by_out_degree If `true`, output is normed by in-degree of
 *                                the output node
 * @param[in]  stream           cuda stream
 * @{
 */
void agg_hg_basis_n2n_pre_fwd(float* out,
                              const float* in,
                              const float* weights_comb,
                              size_t dim,
                              int n_bases,
                              const csc_hg_s32_t& graph,
                              bool concat_own,
                              bool norm_by_out_degree,
                              const cuda::stream& stream);
void agg_hg_basis_n2n_pre_fwd(float* out,
                              const float* in,
                              const float* weights_comb,
                              size_t dim,
                              int n_bases,
                              const csc_hg_s64_t& graph,
                              bool concat_own,
                              bool norm_by_out_degree,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a heterogenous aggregation (RGCN-like)
 *        potentially using a basis decomposition (agg_hg_basis)
 *        operating on graph in a node-to-node reduction (n2n)
 *
 * @note This is based on basis decomposition and assumes that input node
 *       features are concatenated by basis matrix features, i.e. the basis
 *       matrix multiplication is performed before aggregation in forward.
 *       `out_dim`/`dim` is defined as in forward, see the forward docs.
 *
 * @param[out] din             the output gradients on input features.
 *                             [on device] [dim = `graph.n_dst_nodes x dim`].
 *                             This may be `nullptr` if gradients on input nodes
 *                             should be ignored.
 * @param[out] dw              the output weight gradients. [on device]
 *                             [dim = `graph.n_edge_types x n_bases`].
 *                             This may be `nullptr` if the weights were not
 *                             used in forward pass, or gradients on weights
 *                             should be ignored.
 * @param[in]  dout            the input gradients of output features.
 *                             [on device] [dim = `graph.n_dst_nodes x out_dim`]
 * @param[in]  in              the input embeddings used in forward. [on device]
 *                             [dim = `graph.n_dst_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `graph.n_edge_types x n_bases`].
 *                             This may be `nullptr`,
 *                             in which case `n_bases` is ignored.
 * @param[in]  dim             dimensionality of the embeddings. See `forward`.
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`.
 * @param[in]  graph           the input graph.
 * @param[in]  concat_own      See `concat_own` in forward.
 * @param[in]  norm_by_out_degree See `norm_by_out_degree` in forward.
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_hg_basis_n2n_pre_bwd(float* din,
                              float* dw,
                              const float* dout,
                              const float* in,
                              const float* weights_comb,
                              size_t dim,
                              int n_bases,
                              const csc_hg_s32_t& graph,
                              bool concat_own,
                              bool norm_by_out_degree,
                              const cuda::stream& stream);
void agg_hg_basis_n2n_pre_bwd(float* din,
                              float* dw,
                              const float* dout,
                              const float* in,
                              const float* weights_comb,
                              size_t dim,
                              int n_bases,
                              const csc_hg_s64_t& graph,
                              bool concat_own,
                              bool norm_by_out_degree,
                              const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for aggregation layer (RGCN-like) using the
 *        big vertex implementation
 *
 * @note This is based on basis decomposition and assumes that input node
 *       features are concatenated by basis matrix features, i.e. the basis
 *       matrix multiplication is performed before aggregation.
 *       In particular, we define `dim` and `out_dim` as follows:
 *       if (weights_comb != nullptr && !concat_own) dim = n_bases x out_dim
 *       if (weights_comb != nullptr && concat_own)  dim = (n_bases + 1) x out_dim
 *       if (weights_comb == nullptr && !concat_own) dim = graph.n_edge_types x out_dim
 *       if (weights_comb == nullptr && concat_own)  dim = (graph.n_edge_types + 1) x out_dim
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x out_dim`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `graph.n_edge_types x n_bases`].
 *                             This may be `nullptr` if the combination will be
 *                             included in basis matrix multiplication.
 *                             In this case, `n_bases` is ignored.
 * @param[in]  dim             dimensionality of the embeddings.
 *                             See above how it is defined w.r.t. `out_dim`.
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  graph           the input graph.
 * @param[in]  concat_own      If `true`, we assume that node features were
 *                             concatenated to the input embeddings after all
 *                             basis features. In this case, these features
 *                             are aggregated into the output features
 *                             with a weight of `1`.
 * @param[in]  norm_by_out_degree If `true`, output is normed by in-degree of
 *                                the output node
 * @param[in]  workspace        scratch buffer. [on device]
 *                              [len = `workspace_size`] Pass a `nullptr` in
 *                              order to get to know its size.
 * @param[inout] workspace_size workspace size in bytes. If `workspace` is a
 *                              `nullptr`, then this will be computed and the
 *                              caller is expected to allocate the workspace
 *                              buffer of this size.
 * @param[in]  implementation   select implementation type.
 * @param[in]  stream           cuda stream
 * @{
 */
void agg_hg_basis_n2n_pre_big_vertex_fwd(float* out,
                                         const float* in,
                                         const float* weights_comb,
                                         size_t dim,
                                         int n_bases,
                                         const csc_hg_s32_t& graph,
                                         bool concat_own,
                                         bool norm_by_out_degree,
                                         int64_t* workspace,
                                         int64_t& workspace_size,
                                         AggHgImplementationT implementation,
                                         const cuda::stream& stream);
void agg_hg_basis_n2n_pre_big_vertex_fwd(float* out,
                                         const float* in,
                                         const float* weights_comb,
                                         size_t dim,
                                         int n_bases,
                                         const csc_hg_s64_t& graph,
                                         bool concat_own,
                                         bool norm_by_out_degree,
                                         int64_t* workspace,
                                         int64_t& workspace_size,
                                         AggHgImplementationT implementation,
                                         const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a heterogenous aggregation (RGCN-like)
 *        potentially using a basis decomposition (agg_hg_basis)
 *        operating on a graph in a node-to-node
 *        reduction (n2n)
 *
 * @note This is based on basis decomposition and assumes that output node
 *       features are concatenated by basis matrix features, i.e. the basis
 *       matrix multiplication is performed after aggregation.
 *       We will refer to the output dimension as `out_dim`, and it is defined
 *       as follows:
 *       if (weights_comb != nullptr && !concat_own) out_dim = n_bases x dim
 *       if (weights_comb != nullptr && concat_own)  out_dim = (n_bases + 1) x dim
 *       if (weights_comb == nullptr && !concat_own) out_dim = graph.n_edge_types x dim
 *       if (weights_comb == nullptr && concat_own)  out_dim = (graph.n_edge_types + 1) x dim
 *
 * @param[out] out             the output embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x out_dim`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `graph.n_dst_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `graph.n_edge_types x n_bases`].
 *                             This may be `nullptr`,
 *                             in which case `n_bases` is ignored.
 * @param[in]  dim             dimensionality of the input embeddings.
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  graph           the input graph.
 * @param[in]  concat_own      If set, the output node's features are concatenated
 *                             to the output such that the output dimension
 *                             is larger by `dim`.
 * @param[in]  norm_by_out_degree If `true`, output is normed by in-degree of
 *                                the output node
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_hg_basis_n2n_post_fwd(float* out,
                               const float* in,
                               const float* weights_comb,
                               size_t dim,
                               int n_bases,
                               const csc_hg_s32_t& graph,
                               bool concat_own,
                               bool norm_by_out_degree,
                               const cuda::stream& stream);
void agg_hg_basis_n2n_post_fwd(float* out,
                               const float* in,
                               const float* weights_comb,
                               size_t dim,
                               int n_bases,
                               const csc_hg_s64_t& graph,
                               bool concat_own,
                               bool norm_by_out_degree,
                               const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a heterogenous aggregation (RGCN-like)
 *        potentially using a basis decomposition (agg_hg_basis)
 *        operating on a graph in a node-to-node
 *        reduction (n2n)
 *
 * @note This is based on basis decomposition and assumes that output node
 *       features are concatenated by basis matrix features, i.e. the basis
 *       matrix multiplication is performed after aggregation in forward.
 *       `out_dim` is defined as in forward, see the forward docs.
 *
 * @param[out] din             the output gradients on input features. [on device]
 *                             [dim = `graph.n_dst_nodes x dim`].
 *                             This may be `nullptr` if gradients on input nodes
 *                             should be ignored.
 * @param[out] dw              the output weight gradients. [on device]
 *                             [dim = `graph.n_edge_types x n_bases`].
 *                             This may be `nullptr` if the weights were not
 *                             used in forward pass, or gradients on weights
 *                             should be ignored.
 * @param[in]  dout            the input gradients of output features [on device]
 *                             [dim = `graph.n_dst_nodes x out_dim`]
 * @param[in]  in              the input embeddings from forward. [on device]
 *                             [dim = `graph.n_dst_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `graph.n_edge_types x n_bases`].
 *                             This may be `nullptr`,
 *                             in which case `n_bases` is ignored.
 * @param[in]  dim             dimensionality of the input embeddings.
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  graph           the input graph.
 * @param[in]  concat_own      See `concat_own` in forward.
 * @param[in]  norm_by_out_degree See `norm_by_out_degree` in forward.
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_hg_basis_n2n_post_bwd(float* din,
                               float* dw,
                               const float* dout,
                               const float* in,
                               const float* weights_comb,
                               size_t dim,
                               int n_bases,
                               const csc_hg_s32_t& graph,
                               bool concat_own,
                               bool norm_by_out_degree,
                               const cuda::stream& stream);
void agg_hg_basis_n2n_post_bwd(float* din,
                               float* dw,
                               const float* dout,
                               const float* in,
                               const float* weights_comb,
                               size_t dim,
                               int n_bases,
                               const csc_hg_s64_t& graph,
                               bool concat_own,
                               bool norm_by_out_degree,
                               const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
