/*!
 *  Copyright (c) 2019 by Contributors
 * \file graph/nodeflow.cc
 * \brief DGL NodeFlow related functions.
 */

#include <dgl/immutable_graph.h>
#include <dgl/nodeflow.h>

#include <string.h>

#include "../c_api_common.h"

using dgl::runtime::DGLArgs;
using dgl::runtime::DGLArgValue;
using dgl::runtime::DGLRetValue;
using dgl::runtime::PackedFunc;

namespace dgl {

std::vector<IdArray> GetNodeFlowSlice(const ImmutableGraph &graph, const std::string &fmt,
                                      size_t layer0_size, size_t layer1_start,
                                      size_t layer1_end, bool remap) {
  CHECK_GE(layer1_start, layer0_size);
  if (fmt == std::string("csr")) {
    dgl_id_t first_vid = layer1_start - layer0_size;
    auto csr = aten::CSRSliceRows(graph.GetInCSR()->ToCSRMatrix(), layer1_start, layer1_end);
    if (remap) {
      dgl_id_t *eid_data = static_cast<dgl_id_t*>(csr.data->data);
      const dgl_id_t first_eid = eid_data[0];
      IdArray new_indices = aten::Sub(csr.indices, first_vid);
      IdArray new_data = aten::Sub(csr.data, first_eid);
      return {csr.indptr, new_indices, new_data};
    } else {
      return {csr.indptr, csr.indices, csr.data};
    }
  } else if (fmt == std::string("coo")) {
    auto csr = graph.GetInCSR()->ToCSRMatrix();
    const dgl_id_t* indptr = static_cast<dgl_id_t*>(csr.indptr->data);
    const dgl_id_t* indices = static_cast<dgl_id_t*>(csr.indices->data);
    const dgl_id_t* edge_ids = static_cast<dgl_id_t*>(csr.data->data);
    int64_t nnz = indptr[layer1_end] - indptr[layer1_start];
    IdArray idx = aten::NewIdArray(2 * nnz);
    IdArray eid = aten::NewIdArray(nnz);
    int64_t *idx_data = static_cast<int64_t*>(idx->data);
    dgl_id_t *eid_data = static_cast<dgl_id_t*>(eid->data);
    size_t num_edges = 0;
    for (size_t i = layer1_start; i < layer1_end; i++) {
      for (dgl_id_t j = indptr[i]; j < indptr[i + 1]; j++) {
        // These nodes are all in a layer. We need to remap them to the node id
        // local to the layer.
        idx_data[num_edges] = remap ? i - layer1_start : i;
        num_edges++;
      }
    }
    CHECK_EQ(num_edges, nnz);
    if (remap) {
      size_t edge_start = indptr[layer1_start];
      dgl_id_t first_eid = edge_ids[edge_start];
      dgl_id_t first_vid = layer1_start - layer0_size;
      for (int64_t i = 0; i < nnz; i++) {
        CHECK_GE(indices[edge_start + i], first_vid);
        idx_data[nnz + i] = indices[edge_start + i] - first_vid;
        eid_data[i] = edge_ids[edge_start + i] - first_eid;
      }
    } else {
      std::copy(indices + indptr[layer1_start],
                indices + indptr[layer1_end], idx_data + nnz);
      std::copy(edge_ids + indptr[layer1_start],
                edge_ids + indptr[layer1_end], eid_data);
    }
    return std::vector<IdArray>{idx, eid};
  } else {
    LOG(FATAL) << "unsupported adjacency matrix format";
    return {};
  }
}

DGL_REGISTER_GLOBAL("nodeflow._CAPI_NodeFlowGetBlockAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    std::string format = args[1];
    int64_t layer0_size = args[2];
    int64_t start = args[3];
    int64_t end = args[4];
    const bool remap = args[5];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const ImmutableGraph* gptr = dynamic_cast<const ImmutableGraph*>(ptr);
    auto res = GetNodeFlowSlice(*gptr, format, layer0_size, start, end, remap);
    *rv = ConvertNDArrayVectorToPackedFunc(res);
  });

}  // namespace dgl