/*
 * 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 <cuda_runtime.h>

#include <cstdint>

namespace cugraph::ops::dimenet {

/**
 * @brief Computes the forward pass for dimenet RBF features
 *
 * @note N_RADIAL and N_SPHERICAL are constants for now
 *       (equal to 6 and 7 respectively).
 *
 * @param[out] rbf                  the output rbf embeddings incl. envelope. [on device]
 *                                  [dim = `n x N_RADIAL`]
 * @param[out] sbf_rad              the output radial sbf embeddings incl. envelope. [on device]
 *                                  [dim = `n x N_RADIAL x N_SPHERICAL`]
 *                                  [dim = `n x N_RADIAL x N_SPHERICAL`]
 * @param[in]  x                    input distance vector per edge
 *                                  (including offsets) [on device]
 *                                  [dim = `n x 1`]
 * @param[in]  w                    input frequencies [on device] [dim = `N_RADIAL`]
 * @param[in]  n                    number of edges
 * @param[in]  stream               CUDA Stream context
 * @{
 */
void radial_basis_fwd(
  float* rbf, float* sbf_rad, const float* x, const float* w, int32_t n, cudaStream_t stream);
void radial_basis_fwd(
  double* rbf, double* sbf_rad, const double* x, const double* w, int32_t n, cudaStream_t stream);
/** @} */

/**
 * @brief Computes the backward pass for dimenet RBF features
 *
 * @note N_RADIAL and N_SPHERICAL are constants for now
 *       (equal to 6 and 7 respectively).
 *
 * @param[out] grad_x                   output gradient on distance [on device]
 *                                      [dim = `n x 1`]
 * @param[out] grad_w                   output gradient on frequencies. [on device]
 *                                      [dim = `N_RADIAL`]
 * @param[in]  grad_rbf                 the input gradient on rbf features incl. envelope.
 *                                      [on device]
 *                                      [dim = `n x N_RADIAL`]
 * @param[in]  grad_sbf_rad             the input gradient on radial sbf features incl. envelope.
 *                                      [on device]
 *                                      [dim = `n x N_RADIAL x N_SPHERICAL`]
 * @param[in]  x                        input distance vector per edge
 *                                      (including offsets) [on device]
 *                                      [dim = `n x 1`]
 * @param[in]  w                        input frequencies [on device] [dim = `N_RADIAL`]
 * @param[in]  n                        number of edges
 * @param[in]  stream                   CUDA Stream context
 * @{
 */

void radial_basis_bwd(float* grad_x,
                      float* grad_w,
                      const float* grad_rbf,
                      const float* grad_sbf_rad,
                      const float* x,
                      const float* w,
                      int32_t n,
                      cudaStream_t stream);
void radial_basis_bwd(double* grad_x,
                      double* grad_w,
                      const double* grad_rbf,
                      const double* grad_sbf_rad,
                      const double* x,
                      const double* w,
                      int32_t n,
                      cudaStream_t stream);

/** @} */

/**
 * @brief Computes the second backward pass for dimenet RBF features
 *
 * @note N_RADIAL and N_SPHERICAL are constants for now
 *       (equal to 6 and 7 respectively).
 *
 * @param[out] grad_grad_rbf            output gradient on grad_rbf [on device]
 *                                      [dim = `n x N_RADIAL`]
 * @param[out] grad_grad_sbf_rad        output gradient on grad_sbf_rad [on device]
 *                                      [dim = `n x N_RADIAL x N_SPHERICAL]
 * @param[out] grad_w                   output gradient on frequencies. [on device]
 *                                      [dim = `N_RADIAL`]
 * @param[out] grad_grad_x              the input gradient on grad_x [on device]
 *                                      [dim = `n x 1`]
 * @param[out] grad_grad_w              the input gradient on grad_w. [on device]
 *                                      [dim = `N_RADIAL`]
 * @param[in]  grad_rbf                 the input gradient on rbf features incl. envelope.
 *                                      [on device]
 *                                      [dim = `n x N_RADIAL`]
 * @param[in]  grad_sbf_rad             the input gradient on radial sbf features incl. envelope.
 *                                      [on device]
 *                                      [dim = `n x N_RADIAL x N_SPHERICAL`]
 * @param[in]  x                        input distance vector per edge
 *                                      (including offsets) [on device]
 *                                      [dim = `n x 1`]
 * @param[in]  w                        input frequencies [on device] [dim = `N_RADIAL`]
 * @param[in]  n                        number of edges
 * @param[in]  stream                   CUDA Stream context
 * @{
 */

void radial_basis_bwd_bwd(double* grad_grad_rbf,
                          double* grad_grad_sbf_rad,
                          double* grad_w,
                          const double* grad_grad_x,
                          const double* grad_grad_w,
                          const double* grad_rbf,
                          const double* grad_sbf_rad,
                          const double* x,
                          const double* w,
                          int32_t n,
                          cudaStream_t stream);
void radial_basis_bwd_bwd(float* grad_grad_rbf,
                          float* grad_grad_sbf_rad,
                          float* grad_w,
                          const float* grad_grad_x,
                          const float* grad_grad_w,
                          const float* grad_rbf,
                          const float* grad_sbf_rad,
                          const float* x,
                          const float* w,
                          int32_t n,
                          cudaStream_t stream);
/** @} */

}  // namespace cugraph::ops::dimenet
