/*
 * Copyright (c) 2020-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/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 heterogenous aggregation (RGCN-like)
 *        potentially using a basis decomposition (agg_hg_basis)
 *        operatating 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. See `ld`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `mfg.n_out_nodes x ld`]
 * @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]  ld              leading dimension for `out` array. It must be
 *                             at least `n_bases x dim` or
 *                             `mfg.n_edge_types x dim` if `ignore_self`,
 *                             otherwise it must be at least
 *                             `(n_bases + 1) x dim` or
 *                             `(mfg.n_edge_types + 1) x dim`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG. It must be initialized.
 * @param[in]  ignore_self     ignore the self node in the aggregation.
 *                             If `false`, the output node's input features
 *                             will be concatenated to the output features,
 *                             thus `ld` needs to be at least
 *                             `(n_bases + 1) x dim` or
 *                             `(mfg.n_edge_types + 1) x 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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_ellpack_hg_s32_t& mfg,
                                   bool ignore_self,
                                   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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_ellpack_hg_s64_t& mfg,
                                   bool ignore_self,
                                   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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_csr_hg_s32_t& mfg,
                                   bool ignore_self,
                                   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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_csr_hg_s64_t& mfg,
                                   bool ignore_self,
                                   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)
 *        operatating 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.
 *       See `ld`.
 *
 * @param[out] din             the output gradients on input features.
 *                             [on device] [dim = `mfg.n_in_nodes x dim`]
 * @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 ld`]
 * @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]  ld              leading dimension for `dout` array. It must be
 *                             at least `n_bases x dim` or
 *                             `mfg.n_edge_types x dim` if `ignore_self`,
 *                             otherwise it must be at least
 *                             `(n_bases + 1) x dim` or
 *                             `(mfg.n_edge_types + 1) x dim`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`.
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG. It must be initialized.
 * @param[in]  ignore_self     ignore the self node in the aggregation.
 *                             If `false`, it is expected that `ld` is at least
 *                             `(n_bases + 1) x dim` or
 *                             `(mfg.n_edge_types + 1) x dim`.
 *                             See `ignore_self` in forward.
 * @param[in]  norm_by_out_degree Must be set to same value as in forward.
 *                                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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_ellpack_hg_s32_t& mfg,
                                   bool ignore_self,
                                   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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_ellpack_hg_s64_t& mfg,
                                   bool ignore_self,
                                   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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_csr_hg_s32_t& mfg,
                                   bool ignore_self,
                                   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,
                                   size_t ld,
                                   int n_bases,
                                   const mfg_csr_hg_s64_t& mfg,
                                   bool ignore_self,
                                   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)
 *        operatating 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. See `dim`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `mfg.n_out_nodes x ld`]
 * @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. It must be
 *                             a multiple of `n_bases` if weights are used,
 *                             or a multiple of `mfg.n_edge_types` otherwise.
 *                             The embeddings must be ordered such that the
 *                             first `dim / n_bases` elements correspond to
 *                             features of the first base matrix and so on.
 * @param[in]  ld              leading dimension for `out` array. It must be
 *                             at least `dim/n_bases` if weights are being used.
 *                             Otherwise it must be at least `dim/mfg.n_edge_types`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG. It must be initialized.
 * @param[in]  type_conv       Only relevant if `!ignore_self`: how to convert
 *                             the node type to an edge type for a self-loop.
 *                             Pass a `nullptr` if you want to ignore this.
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_ellpack_hg_s32_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_ellpack_hg_s64_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_csr_hg_s32_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_csr_hg_s64_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  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)
 *        operatating 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.
 *       See `dim`.
 *
 * @param[out] din             the output gradients on input features.
 *                             [on device] [dim = `mfg.n_in_nodes x dim`]
 * @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 ld`]
 * @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. It must be
 *                             a multiple of `n_bases` if weights are used,
 *                             or a multiple of `mfg.n_edge_types` otherwise.
 *                             The embeddings must be ordered such that the
 *                             first `dim / n_bases` elements correspond to
 *                             features of the first base matrix and so on.
 * @param[in]  ld              leading dimension for `out` array. It must be
 *                             at least `dim/n_bases` if weights are being used.
 *                             Otherwise it must be at least `dim/mfg.n_edge_types`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`.
 * @param[in]  mfg             the input (ELLPACK/CSR) MFG. It must be initialized.
 * @param[in]  type_conv       Only relevant if `!ignore_self`: how to convert
 *                             the node type to an edge type for a self-loop.
 *                             Pass a `nullptr` if you want to ignore this.
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @param[in]  norm_by_out_degree Must be set to same value as in forward.
 *                                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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_ellpack_hg_s32_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_ellpack_hg_s64_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_csr_hg_s32_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  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,
                                  size_t ld,
                                  int n_bases,
                                  const mfg_csr_hg_s64_t& mfg,
                                  const self_loop_type_conv_t* type_conv,
                                  bool ignore_self,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
