/*
 * 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/operators/common.hpp>
#include <cugraph-ops/test/graph/format_mfg.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 message flow graphs (mfg) 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 = mfg.n_edge_types x dim
 *       if (weights_comb == nullptr && concat_own)  out_dim = (mfg.n_edge_types + 1) x dim
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `mfg.n_out_nodes x out_dim`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `mfg.n_in_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `mfg.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.
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG.
 * @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
 *                                per-edge-type-in-degree of the output node
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_hg_basis_mfg_n2n_post_fwd(float* out,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_ellpack_hg_s32_t& graph,
                                   bool concat_own,
                                   bool norm_by_out_degree,
                                   const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_post_fwd(float* out,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_ellpack_hg_s64_t& graph,
                                   bool concat_own,
                                   bool norm_by_out_degree,
                                   const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_post_fwd(float* out,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_csc_hg_s32_t& graph,
                                   bool concat_own,
                                   bool norm_by_out_degree,
                                   const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_post_fwd(float* out,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_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 message flow graphs (mfg) 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 = `mfg.n_in_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 = `mfg.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 = `mfg.n_out_nodes x out_dim`]
 * @param[in]  in              the input embeddings used in forward. [on device]
 *                             [dim = `mfg.n_in_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `mfg.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]  mfg             the input (ELLPACK/CSR) MFG.
 * @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_mfg_n2n_post_bwd(float* din,
                                   float* dw,
                                   const float* dout,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_ellpack_hg_s32_t& graph,
                                   bool concat_own,
                                   bool norm_by_out_degree,
                                   const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_post_bwd(float* din,
                                   float* dw,
                                   const float* dout,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_ellpack_hg_s64_t& graph,
                                   bool concat_own,
                                   bool norm_by_out_degree,
                                   const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_post_bwd(float* din,
                                   float* dw,
                                   const float* dout,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_csc_hg_s32_t& graph,
                                   bool concat_own,
                                   bool norm_by_out_degree,
                                   const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_post_bwd(float* din,
                                   float* dw,
                                   const float* dout,
                                   const float* in,
                                   const float* weights_comb,
                                   size_t dim,
                                   int n_bases,
                                   const mfg_csc_hg_s64_t& graph,
                                   bool concat_own,
                                   bool norm_by_out_degree,
                                   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 message flow graphs (mfg) 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 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 = mfg.n_edge_types x out_dim
 *       if (weights_comb == nullptr && concat_own)  dim = (mfg.n_edge_types + 1) x out_dim
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `mfg.n_out_nodes x out_dim`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `mfg.n_in_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `mfg.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]  mfg             the input (ELLPACK/CSR) MFG.
 * @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
 *                                per-edge-type-in-degree of the output node
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_hg_basis_mfg_n2n_pre_fwd(float* out,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_ellpack_hg_s32_t& graph,
                                  bool concat_own,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_pre_fwd(float* out,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_ellpack_hg_s64_t& graph,
                                  bool concat_own,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_pre_fwd(float* out,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_csc_hg_s32_t& graph,
                                  bool concat_own,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_pre_fwd(float* out,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_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 message flow graphs (mfg) 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 = `mfg.n_in_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 = `mfg.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 = `mfg.n_out_nodes x out_dim`]
 * @param[in]  in              the input embeddings used in forward. [on device]
 *                             [dim = `mfg.n_in_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `mfg.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]  mfg             the input (ELLPACK/CSR) MFG.
 * @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_mfg_n2n_pre_bwd(float* din,
                                  float* dw,
                                  const float* dout,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_ellpack_hg_s32_t& graph,
                                  bool concat_own,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_pre_bwd(float* din,
                                  float* dw,
                                  const float* dout,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_ellpack_hg_s64_t& graph,
                                  bool concat_own,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_pre_bwd(float* din,
                                  float* dw,
                                  const float* dout,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_csc_hg_s32_t& graph,
                                  bool concat_own,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_mfg_n2n_pre_bwd(float* din,
                                  float* dw,
                                  const float* dout,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  int n_bases,
                                  const mfg_csc_hg_s64_t& graph,
                                  bool concat_own,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
