/*
 * 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 {

enum class AggHgImplementationT {
  kDefaultT,
  kBigVertexImplementationT,
  kBigVertexLessMemoryImplementationT
};

/**
 * @brief Computes the forward pass for a heterogenous aggregation (RGCN-like)
 *        potentially using a basis decomposition (agg_hg_basis)
 *        operatating on full graphs (fg) 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. See `dim`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x ld`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `fg.n_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `fg.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.
 *                             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 `fg.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/fg.n_edge_types`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  fg              the input CSR full graph
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @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_fg_n2n_pre_fwd(float* out,
                                 const float* in,
                                 const float* weights_comb,
                                 size_t dim,
                                 size_t ld,
                                 int n_bases,
                                 const fg_csr_hg_s32_t& fg,
                                 bool ignore_self,
                                 bool norm_by_out_degree,
                                 const cuda::stream& stream);
void agg_hg_basis_fg_n2n_pre_fwd(float* out,
                                 const float* in,
                                 const float* weights_comb,
                                 size_t dim,
                                 size_t ld,
                                 int n_bases,
                                 const fg_csr_hg_s64_t& fg,
                                 bool ignore_self,
                                 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. See `dim`.
 *
 * @param[out] out             the output aggregated embeddings. [on device]
 *                             [dim = `fg.n_nodes x ld`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `fg.n_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `fg.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 `fg.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/fg.n_edge_types`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  fg              the input CSR full graph
 * @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 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_fg_n2n_pre_big_vertex_fwd(float* out,
                                            const float* in,
                                            const float* weights_comb,
                                            size_t dim,
                                            size_t ld,
                                            int n_bases,
                                            const fg_csr_hg_s32_t& fg,
                                            const self_loop_type_conv_t* type_conv,
                                            bool ignore_self,
                                            bool norm_by_out_degree,
                                            int64_t* workspace,
                                            int64_t& workspace_size,
                                            AggHgImplementationT implementation,
                                            const cuda::stream& stream);
void agg_hg_basis_fg_n2n_pre_big_vertex_fwd(float* out,
                                            const float* in,
                                            const float* weights_comb,
                                            size_t dim,
                                            size_t ld,
                                            int n_bases,
                                            const fg_csr_hg_s64_t& fg,
                                            const self_loop_type_conv_t* type_conv,
                                            bool ignore_self,
                                            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)
 *        operatating on full graphs (fg) 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 embeddings. [on device]
 *                             [dim = `fg.n_nodes x ld`]
 * @param[in]  in              the input embeddings. [on device]
 *                             [dim = `fg.n_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `fg.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 `out` array. It must be
 *                             at least `n_bases x dim` or
 *                             `fg.n_edge_types x dim` if `ignore_self`,
 *                             otherwise it must be at least
 *                             `(n_bases + 1) x dim` or
 *                             `(fg.n_edge_types + 1) x dim`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  fg              the input CSR heterogeneous full graph
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @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_fg_n2n_post_fwd(float* out,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  size_t ld,
                                  int n_bases,
                                  const fg_csr_hg_s32_t& fg,
                                  bool ignore_self,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_fg_n2n_post_fwd(float* out,
                                  const float* in,
                                  const float* weights_comb,
                                  size_t dim,
                                  size_t ld,
                                  int n_bases,
                                  const fg_csr_hg_s64_t& fg,
                                  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 full graphs (fg) 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] din             the output gradients on input features. [on device]
 *                             [dim = `fg.n_nodes x dim`]
 * @param[out] dw              the output weight gradients. [on device]
 *                             [dim = `fg.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 = `fg.n_nodes x ld`]
 * @param[in]  in              the input embeddings from forward. [on device]
 *                             [dim = `fg.n_nodes x dim`]
 * @param[in]  weights_comb    the edge type combination weights. [on device]
 *                             [dim = `fg.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 `out` array. It must be
 *                             at least `n_bases x dim` or
 *                             `fg.n_edge_types x dim` if `ignore_self`,
 *                             otherwise it must be at least
 *                             `(n_bases + 1) x dim` or
 *                             `(fg.n_edge_types + 1) x dim`
 * @param[in]  n_bases         Number of bases used.
 *                             Ignored if `weights_comb == nullptr`
 * @param[in]  fg              the input CSR heterogeneous full graph
 * @param[in]  ignore_self     ignore the self node in the aggregation
 * @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_fg_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 fg_csr_hg_s32_t& fg,
                                  bool ignore_self,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
void agg_hg_basis_fg_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 fg_csr_hg_s64_t& fg,
                                  bool ignore_self,
                                  bool norm_by_out_degree,
                                  const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops::gnn
