/**
 *  Copyright (c) 2020 by Contributors
 * @file dgl/aten/kernel.h
 * @brief Sparse matrix operators.
 */
#ifndef DGL_KERNEL_H_
#define DGL_KERNEL_H_

#include <string>
#include <utility>
#include <vector>

#include "./base_heterograph.h"
#include "./bcast.h"
#include "array.h"

namespace dgl {
namespace aten {

/**
 * @brief Generalized Sparse Matrix-Matrix Multiplication.
 * @param op The binary operator, could be `add`, `sub', `mul`, 'div',
 *        `copy_u`, `copy_e'.
 * @param op The reduce operator, could be `sum`, `min`, `max'.
 * @param graph The graph we apply SpMM on.
 * @param ufeat The source node feature.
 * @param efeat The edge feature.
 * @param out The output feature on destination nodes.
 * @param out_aux A list of NDArray's that contains auxiliary information such
 *        as the argmax on source nodes and edges for reduce operators such as
 *        `min` and `max`.
 */
void SpMM(
    const std::string& op, const std::string& reduce, HeteroGraphPtr graph,
    NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);

/**
 * @brief Generalized Sampled Dense-Dense Matrix Multiplication.
 * @param op The binary operator, could be `add`, `sub', `mul`, 'div',
 *        `dot`, `copy_u`, `copy_e'.
 * @param graph The graph we apply SpMM on.
 * @param ufeat The source node feature.
 * @param vfeat The destination node feature.
 * @param out The output feature on edge.
 */
void SDDMM(
    const std::string& op, HeteroGraphPtr graph, NDArray ufeat, NDArray efeat,
    NDArray out);

/**
 * @brief Sparse-sparse matrix multiplication.
 *
 * The sparse matrices must have scalar weights (i.e. \a A_weights and \a
 * B_weights are 1D vectors.)
 */
std::pair<CSRMatrix, NDArray> CSRMM(
    CSRMatrix A, NDArray A_weights, CSRMatrix B, NDArray B_weights);

/**
 * @brief Summing up a list of sparse matrices.
 *
 * The sparse matrices must have scalar weights (i.e. the arrays in \a A_weights
 * are 1D vectors.)
 */
std::pair<CSRMatrix, NDArray> CSRSum(
    const std::vector<CSRMatrix>& A, const std::vector<NDArray>& A_weights);

}  // namespace aten
}  // namespace dgl

#endif  // DGL_KERNEL_H_