/*
 * 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/mma.hpp>

#include <cuda_runtime.h>

#include <cstdint>

namespace cugraph::ops::dimenet {

// TODO(mjoux): unused in the API, but only used for reference for how to
//              use triplet indexes if comparing to baseline
// template <typename DataT, typename IdxT>
// struct AuxiliaryStorage {
//   const IdxT *triplet_offsets, *triplet_idx = nullptr;
//   DataT *sbf_dump, *mid_dump, *feat_dump;
// };

void agg_edge_to_edge_fwd(float* out,
                          int64_t n_edges,
                          const float* e_vec,
                          const float* e_rbf,
                          const float* e_in_ft,
                          size_t dim_feat,
                          const float* sbf_weights,
                          size_t dim_w_mid,
                          const int64_t* coo_idx,
                          const int64_t* dst_offsets,
                          const int64_t* dst_e_idx,
                          MMAOpT op,
                          cudaStream_t stream);
void agg_edge_to_edge_fwd(double* out,
                          int64_t n_edges,
                          const double* e_vec,
                          const double* e_rbf,
                          const double* e_in_ft,
                          size_t dim_feat,
                          const double* sbf_weights,
                          size_t dim_w_mid,
                          const int64_t* coo_idx,
                          const int64_t* dst_offsets,
                          const int64_t* dst_e_idx,
                          MMAOpT op,
                          cudaStream_t stream);

void agg_edge_to_edge_bwd(float* grad_e_rbf,
                          float* grad_e_in_ft,
                          float* grad_sbf_weights,
                          int64_t n_edges,
                          const float* e_vec,
                          const float* e_rbf,
                          const float* e_in_ft,
                          const float* grad_e_out_ft,
                          size_t dim_feat,
                          const float* sbf_weights,
                          size_t dim_w_mid,
                          const int64_t* coo_idx,
                          const int64_t* src_offsets,
                          MMAOpT op,
                          cudaStream_t stream);
void agg_edge_to_edge_bwd(double* grad_e_rbf,
                          double* grad_e_in_ft,
                          double* grad_sbf_weights,
                          int64_t n_edges,
                          const double* e_vec,
                          const double* e_rbf,
                          const double* e_in_ft,
                          const double* grad_e_out_ft,
                          size_t dim_feat,
                          const double* sbf_weights,
                          size_t dim_w_mid,
                          const int64_t* coo_idx,
                          const int64_t* src_offsets,
                          MMAOpT op,
                          cudaStream_t stream);

void agg_edge_to_edge_bwd2_grad(float* grad_grad_e_out_ft,
                                int64_t n_edges,
                                const float* e_vec,
                                const float* e_rbf,
                                const float* grad_grad_e_rbf,
                                const float* e_in_ft,
                                const float* grad_grad_e_in_ft,
                                size_t dim_feat,
                                const float* sbf_weights,
                                size_t dim_w_mid,
                                const int64_t* coo_idx,
                                const int64_t* dst_offsets,
                                const int64_t* dst_e_idx,
                                MMAOpT op,
                                cudaStream_t stream);
void agg_edge_to_edge_bwd2_grad(double* grad_grad_e_out_ft,
                                int64_t n_edges,
                                const double* e_vec,
                                const double* e_rbf,
                                const double* grad_grad_e_rbf,
                                const double* e_in_ft,
                                const double* grad_grad_e_in_ft,
                                size_t dim_feat,
                                const double* sbf_weights,
                                size_t dim_w_mid,
                                const int64_t* coo_idx,
                                const int64_t* dst_offsets,
                                const int64_t* dst_e_idx,
                                MMAOpT op,
                                cudaStream_t stream);

void agg_edge_to_edge_bwd2_main(float* grad_rbf_grad_e_in_ft,
                                float* grad_e_in_ft_grad_e_rbf,
                                float* grad_sbf_weights,
                                int64_t n_edges,
                                const float* e_vec,
                                const float* e_rbf,
                                const float* grad_grad_e_rbf,
                                const float* e_in_ft,
                                const float* grad_grad_e_in_ft,
                                const float* grad_e_out_ft,
                                size_t dim_feat,
                                const float* sbf_weights,
                                size_t dim_w_mid,
                                const int64_t* coo_idx,
                                const int64_t* src_offsets,
                                MMAOpT op,
                                cudaStream_t stream);
void agg_edge_to_edge_bwd2_main(double* grad_rbf_grad_e_in_ft,
                                double* grad_e_in_ft_grad_e_rbf,
                                double* grad_sbf_weights,
                                int64_t n_edges,
                                const double* e_vec,
                                const double* e_rbf,
                                const double* grad_grad_e_rbf,
                                const double* e_in_ft,
                                const double* grad_grad_e_in_ft,
                                const double* grad_e_out_ft,
                                size_t dim_feat,
                                const double* sbf_weights,
                                size_t dim_w_mid,
                                const int64_t* coo_idx,
                                const int64_t* src_offsets,
                                MMAOpT op,
                                cudaStream_t stream);

}  // namespace cugraph::ops::dimenet
