/*!
 *  Copyright (c) 2018 by Contributors
 * \file dgl/scheduler.h
 * \brief Operations on graph index.
 */
#ifndef DGL_SCHEDULER_H_
#define DGL_SCHEDULER_H_

#include <vector>
#include "runtime/ndarray.h"

namespace dgl {

typedef dgl::runtime::NDArray IdArray;

namespace sched {

/*!
 * \brief Generate degree bucketing schedule
 * \tparam IdType Graph's index data type, can be int32_t or int64_t
 * \param msg_ids The edge id for each message
 * \param vids The destination vertex for each message
 * \param recv_ids The recv nodes (for checking zero degree nodes)
 * \note If there are multiple messages going into the same destination vertex, then
 *       there will be multiple copies of the destination vertex in vids
 * \return a vector of 5 IdArrays for degree bucketing. The 5 arrays are:
 *         degrees: degrees for each bucket
 *         nids: destination node ids
 *         nid_section: number of nodes in each bucket (used to split nids)
 *         mids: message ids
 *         mid_section: number of messages in each bucket (used to split mids)
 */
template <class IdType>
std::vector<IdArray> DegreeBucketing(const IdArray& msg_ids, const IdArray& vids,
        const IdArray& recv_ids);

/*!
 * \brief Generate degree bucketing schedule for group_apply edge
 * \tparam IdType Graph's index data type, can be int32_t or int64_t
 * \param uids One end vertex of edge by which edges are grouped
 * \param vids The other end vertex of edge
 * \param eids Edge ids
 * \note This function always generate group_apply schedule based on degrees of
 *       nodes in uids. Therefore, if group_apply by source nodes, then uids
 *       should be source. If group_apply by destination nodes, then uids
 *       should be destination.
 * \return a vector of 5 IdArrays for degree bucketing. The 5 arrays are:
 *         degrees: degrees for each bucket
 *         new_uids: uids reordered by degree bucket
 *         new_vids: vids reordered by degree bucket
 *         new_edis: eids reordered by degree bucket
 *         sections: number of edges in each degree bucket (used to partition
 *                   new_uids, new_vids, and new_eids)
 */
template <class IdType>
std::vector<IdArray> GroupEdgeByNodeDegree(const IdArray& uids,
        const IdArray& vids, const IdArray& eids);

}  // namespace sched

}  // namespace dgl

#endif  // DGL_SCHEDULER_H_