/*
 * 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/graph/format.hpp>
#include <cugraph-ops/operators/common.hpp>

#include <cstdint>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for D-MPNN-like full graph aggregation
 *        (agg_dmpnn_csc) using edge features in an edge-to-edge reduction
 *        scheme (e2e)
 *
 * @note For these operations, if the graph uses edge indices
 *       (`graph.map_csc_to_coo != nullptr`), then these indices must be a partition /
 *       permutation of the edge IDs.
 *
 * @param[out] out_edge_feat   the output edge features [on device]
 *                             output dimension `out_dim` is `dim_edge * 2` if
 *                             `concat_own`, otherwise `dim_edge`.
 * @param[in]  in_edge_feat    the input edge features. [on device]
 *                             [dim = `graph.n_indices x dim_edge`].
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  graph           the input CSR full graph. It must be undirected:
 *                             for every edge, there is a reverse edge present
 *                             in the structure, too.
 * @param[in]  rev_edge_ids    position indices for each edge
 *                             of the corresponding reverse edge,
 *                             i.e. its corresponding position in graph.indices
 * @param[in]  concat_own      If `true`, concatenate an edge's input features
 *                             to the aggregated output features.
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_dmpnn_e2e_fwd(float* out_edge_feat,
                       const float* in_edge_feat,
                       size_t dim_edge,
                       const csc_s32_t& graph,
                       const int32_t* rev_edge_ids,
                       bool concat_own,
                       const cuda::stream& stream);
void agg_dmpnn_e2e_fwd(float* out_edge_feat,
                       const float* in_edge_feat,
                       size_t dim_edge,
                       const csc_s64_t& graph,
                       const int64_t* rev_edge_ids,
                       bool concat_own,
                       const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for D-MPNN-like full graph aggregation
 *        (agg_dmpnn_csc) using edge features in an edge-to-edge reduction
 *        scheme (e2e)
 *
 * @note For these operations, if the graph uses edge indices
 *       (`graph.map_csc_to_coo != nullptr`), then these indices must be a partition /
 *       permutation of the edge IDs.
 *
 * @param[out] d_in_edge_feat  the output gradients on input edge features
 *                             [on device] [dim = `graph.n_indices x dim_edge`].
 * @param[in]  d_out_edge_feat the input gradients on output edge features.
 *                             [on device] [dim = `graph.n_indices x out_dim`].
 *                             For `out_dim`, see parameter `out_edge_feat`
 *                             of the forward pass.
 * @param[in]  dim_edge        dimensionality of the edge features
 * @param[in]  graph           the input CSR full graph. It must be undirected:
 *                             for every edge, there is a reverse edge present
 *                             in the structure, too.
 * @param[in]  rev_edge_ids    position indices for each edge
 *                             of the corresponding reverse edge,
 *                             i.e. its corresponding position in graph.indices
 * @param[in]  concat_own      If `true`, concatenate an edge's input features
 *                             to the aggregated output features.
 * @param[in]  stream          cuda stream
 * @{
 */
void agg_dmpnn_e2e_bwd(float* d_in_edge_feat,
                       const float* d_out_edge_feat,
                       size_t dim_edge,
                       const csc_s32_t& graph,
                       const int32_t* rev_edge_ids,
                       bool concat_own,
                       const cuda::stream& stream);
void agg_dmpnn_e2e_bwd(float* d_in_edge_feat,
                       const float* d_out_edge_feat,
                       size_t dim_edge,
                       const csc_s64_t& graph,
                       const int64_t* rev_edge_ids,
                       bool concat_own,
                       const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
