/*!
 *  Copyright (c) 2018 by Contributors
 * \file graph/traversal.cc
 * \brief Graph traversal implementation
 */
#include <dgl/graph_traversal.h>
#include <dgl/packed_func_ext.h>

#include "../c_api_common.h"

using namespace dgl::runtime;

namespace dgl {
namespace traverse {

DGL_REGISTER_GLOBAL("traversal._CAPI_DGLBFSNodes_v2")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef g = args[0];
      const IdArray src = args[1];
      bool reversed = args[2];
      aten::CSRMatrix csr;
      if (reversed) {
        csr = g.sptr()->GetCSCMatrix(0);
      } else {
        csr = g.sptr()->GetCSRMatrix(0);
      }
      const auto& front = aten::BFSNodesFrontiers(csr, src);
      *rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.sections});
    });

DGL_REGISTER_GLOBAL("traversal._CAPI_DGLBFSEdges_v2")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef g = args[0];
      const IdArray src = args[1];
      bool reversed = args[2];
      aten::CSRMatrix csr;
      if (reversed) {
        csr = g.sptr()->GetCSCMatrix(0);
      } else {
        csr = g.sptr()->GetCSRMatrix(0);
      }

      const auto& front = aten::BFSEdgesFrontiers(csr, src);
      *rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.sections});
    });

DGL_REGISTER_GLOBAL("traversal._CAPI_DGLTopologicalNodes_v2")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef g = args[0];
      bool reversed = args[1];
      aten::CSRMatrix csr;
      if (reversed) {
        csr = g.sptr()->GetCSCMatrix(0);
      } else {
        csr = g.sptr()->GetCSRMatrix(0);
      }

      const auto& front = aten::TopologicalNodesFrontiers(csr);
      *rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.sections});
    });

DGL_REGISTER_GLOBAL("traversal._CAPI_DGLDFSEdges_v2")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef g = args[0];
      const IdArray source = args[1];
      const bool reversed = args[2];
      CHECK(aten::IsValidIdArray(source)) << "Invalid source node id array.";
      aten::CSRMatrix csr;
      if (reversed) {
        csr = g.sptr()->GetCSCMatrix(0);
      } else {
        csr = g.sptr()->GetCSRMatrix(0);
      }
      const auto& front = aten::DGLDFSEdges(csr, source);
      *rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.sections});
    });

DGL_REGISTER_GLOBAL("traversal._CAPI_DGLDFSLabeledEdges_v2")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef g = args[0];
      const IdArray source = args[1];
      const bool reversed = args[2];
      const bool has_reverse_edge = args[3];
      const bool has_nontree_edge = args[4];
      const bool return_labels = args[5];
      aten::CSRMatrix csr;
      if (reversed) {
        csr = g.sptr()->GetCSCMatrix(0);
      } else {
        csr = g.sptr()->GetCSRMatrix(0);
      }

      const auto& front = aten::DGLDFSLabeledEdges(
          csr, source, has_reverse_edge, has_nontree_edge, return_labels);

      if (return_labels) {
        *rv = ConvertNDArrayVectorToPackedFunc(
            {front.ids, front.tags, front.sections});
      } else {
        *rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.sections});
      }
    });

}  // namespace traverse
}  // namespace dgl