/*
 * 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/max_vertex_degree.hpp>
#include <cugraph-ops/test/graph/format_mfg.hpp>
#include <cugraph-ops/test/graph/output_batching.hpp>
#include <cugraph-ops/test/graph/renumber.hpp>
#include <cugraph-ops/test/graph/reverse_graph_mfg.hpp>
#include <cugraph-ops/test/graph/sampling_mfg.hpp>
#include <cugraph-ops/test/graph/uniquify.hpp>
#include <cugraph-ops/utils/allocator.hpp>
#include <cugraph-ops/utils/error.hpp>
#include <cugraph-ops/utils/logger.hpp>
#include <cugraph-ops/utils/nvtx.hpp>

#include <raft/random/rng_state.hpp>

#include <cuda_runtime.h>

#include <algorithm>
#include <cstdint>
#include <memory>
#include <random>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

namespace cugraph::ops {

/**
 * @brief Generates the batch of nodes along with neighborhood sampling, for the
 *        current iteration
 *
 * @code{.cpp}
 * node_batcher<IdxT> b(sizes, batch_size, n, rng, stream);
 * for (int i = 0; i < b.num_batches(); ++i) {
 *   IdxT batch_size, *batch;
 *   std::tie(batch_size, batch) = b.next_batch();
 * }
 * @endcode
 */
template <typename IdxT, bool USE_CSR = false>
struct node_batcher {
  static constexpr bool IS_CSR = USE_CSR;
  using csc_t                  = graph::csc<IdxT>;
  using mfg_t =
    typename std::conditional<USE_CSR, graph::mfg_csc<IdxT>, graph::mfg_ellpack<IdxT>>::type;
  using graph_rev_t = graph::mfg_csc_rev<IdxT>;

  /**
   * @brief Info related to the current batch
   */
  struct __attribute__((aligned(16))) batch_info {
    /** pointer to the current output batch of nodes */
    IdxT* batch{nullptr};

    /** size of the current batch */
    IdxT curr_size{0};
  };  // struct batch_info

  /**
   * @brief All of the nodes and their neighborhood info for a given layer
   *
   * @note None of the pointers here are owned by this class, and all the
   *       pointers are on device, unless mentioned otherwise.
   */
  struct layer_info {
    // handle sub-classes correctly
    // unused for now but may be useful in the future
    layer_info()          = default;
    virtual ~layer_info() = default;

    /**
     * The main MFG (see graph/format.hpp for details).
     * The output nodes are renumbered as per the input nodes of this layer.
     * The input nodes are uniquified and renumbered as per the input nodes
     * of the previous layer.
     * For the first layer, the input nodes simply contain global graph IDs.
     */
    std::unique_ptr<mfg_t> graph;
    /**
     * The reverse MFG (see graph/format.hpp for details).
     * It is only populated if `node_batcher.do_reverse_` is set.
     */
    std::unique_ptr<graph::mfg_csc_rev<IdxT>> graph_rev;
    /**
     * Input nodes of the MFG.
     * In middle layers, same as output nodes from previous layer.
     */
    IdxT* in_nodes{nullptr};
    /** Number of input nodes of the MFG. Only used for calculating it. */
    IdxT* n_in_nodes{nullptr};
    /** Number of indices/edges of the MFG. Only used for calculating it. */
    IdxT* n_indices{nullptr};

    /** sample size: maximum number of neighbors of an output node */
    IdxT sample_size{0};
    /** maximum possible number of output nodes for this layer */
    IdxT max_out_nodes{0};
  };  // struct layer_info

  /**
   * @param[in]    _graph      the graph to be sub-sampled
   * @param[in]    _sizes      list of sample sizes at each layer of network.
   *                           Its length will tell the number of layers in the
   *                           network. Sample sizes must be mentioned starting
   *                           from the first layer, all the way to the output
   *                           layer, in that order only.
   * @param[in]    _batch_size batch size
   * @param[in]    _n          total number of nodes
   * @param[in]    _n_labels   number of labeled nodes (the graph node IDs must
   *                           be ordered such that the first `_n_labels`
   *                           nodes are labeled)
   * @param[inout] _rng        RAFT RngState object
   * @param[in]    _type       sampling type
   * @param[in]    _do_reverse whether to compute a reverse graph
   * @param[in]    _stream     cuda stream
   */
  node_batcher(csc_t& _graph,
               const std::vector<IdxT>& _sizes,
               IdxT _batch_size,
               IdxT _n,
               IdxT _n_labels,
               raft::random::RngState& _rng,
               graph::SamplingAlgoT _type,
               bool _do_reverse,
               const cuda::stream& _stream)
    : graph_(_graph),
      sample_sizes_(_sizes),
      batch_size_(_batch_size),
      n_(_n),
      n_labels_(_n_labels),
      rng_(_rng),
      gen_(_rng.seed),
      sample_type_(_type),
      do_reverse_(_do_reverse),
      stream_(_stream),
      num_batches_((n_labels_ + batch_size_ - 1) / batch_size_),
      info_(_sizes.size())
  {
    for (size_t i = 0; i < sample_sizes_.size(); ++i) {
      info_[i]        = std::make_unique<layer_info>();
      info_[i]->graph = std::make_unique<mfg_t>();
      if (do_reverse_) info_[i]->graph_rev = std::make_unique<graph_rev_t>();
    }
    assign_workspace();
    reset();
  }

  /**
   * Node batcher without graph reversal
   *
   * @param[in]    _graph      the graph to be sub-sampled
   * @param[in]    _sizes      list of sample sizes at each layer of network.
   *                           Its length will tell the number of layers in the
   *                           network. Sample sizes must be mentioned starting
   *                           from the first layer, all the way to the output
   *                           layer, in that order only.
   * @param[in]    _batch_size batch size
   * @param[in]    _n          total number of nodes (all are labeled)
   * @param[inout] _rng        RAFT RngState object
   * @param[in]    _type       sampling type
   * @param[in]    _stream     cuda stream
   */
  node_batcher(csc_t& _graph,
               const std::vector<IdxT>& _sizes,
               IdxT _batch_size,
               IdxT _n,
               raft::random::RngState& _rng,
               graph::SamplingAlgoT _type,
               const cuda::stream& _stream)
    : node_batcher(_graph, _sizes, _batch_size, _n, _n, _rng, _type, false, _stream)
  {
  }

  /** Resets the batch generator (eg: for the next epoch) */
  void reset()
  {
    next_start_[0]   = 0;
    xor_key_         = gen_();
    count_           = 0;
    curr_batch_size_ = 0;
  }

  ~node_batcher()
  {
    auto& h_alloc = utils::host_allocator();
    auto& d_alloc = utils::device_allocator();
    h_alloc.dealloc(next_start_, 0, stream_());
    d_alloc.dealloc(batch_, 0, stream_());
    for (size_t i = 0; i < info_.size(); ++i) {
      dealloc_layer_info_fields(d_alloc, h_alloc, i);
    }
    d_alloc.dealloc(workspace_, 0, stream_());
  }

  /** node batcher cannot be copied because of the unique pointers to layer info */
  node_batcher(const node_batcher&)            = delete;
  node_batcher& operator=(const node_batcher&) = delete;

  /**
   * @brief computes the next batch of nodes
   *
   * @return a pair containing the actual batch size of the output nodes and
   *         device pointer containing the list of output node id's (id's as per
   *         the original graph) in this batch
   *
   * @note the output batch pointer should NOT be `cudaFree`d by the caller
   */
  virtual batch_info next_batch()
  {
    utils::push_range("node_batcher::next_batch");
    // generate next output batch of nodes, reset if all nodes covered already
    if (n_labels_ <= count_ * batch_size_) { reset(); }
    curr_batch_size_ = std::max(IdxT{0}, std::min(n_labels_ - count_ * batch_size_, batch_size_));
    graph::get_next_output_batch(batch_,
                                 next_start_,
                                 curr_batch_size_,
                                 n_labels_,
                                 xor_key_,
                                 workspace_,
                                 workspace_size_,
                                 stream_());
    ++count_;
    // start from output layer to compute the sub-graph for message passing
    IdxT* out_nodes = batch_;
    auto last_id    = static_cast<int>(info_.size() - 1);
    for (auto i = last_id; i >= 0; --i) {
      auto n_out_nodes = i == last_id ? curr_batch_size_ : info_[i + 1]->graph->n_in_nodes;
      // update the #out nodes for each layer too
      // this must be done before sampling/uniquify so that these functions
      // can rely on this number to be set
      info_[i]->graph->n_out_nodes = n_out_nodes;
      graph::uniform_sample(*info_[i]->graph,
                            rng_,
                            graph_,
                            out_nodes,
                            sample_type_,
                            max_degree_,
                            sample_self_,
                            info_[i]->n_indices,
                            workspace_,
                            workspace_size_,
                            stream_());
      graph::uniquify(info_[i]->in_nodes,
                      info_[i]->n_in_nodes,
                      *info_[i]->graph,
                      out_nodes,
                      n_out_nodes,
                      n_ + 1,
                      workspace_,
                      workspace_size_,
                      stream_());
      // uniquify should sync, so we're able to access `n_in_nodes` here
      info_[i]->graph->n_in_nodes = info_[i]->n_in_nodes[0];
      out_nodes                   = info_[i]->in_nodes;
    }
    copy_additional_info();
    renumber_node_ids(curr_batch_size_);
    utils::pop_range();
    return curr_batch();
  }

  /** returns the current batch information */
  batch_info curr_batch() { return batch_info{batch_, curr_batch_size_}; }

  /** total number of batches to be generated */
  [[nodiscard]] IdxT num_batches() const { return num_batches_; }

  /** total number of layers */
  [[nodiscard]] size_t num_layers() const { return info_.size(); }

  /** fetches the given layer's node/sampling information */
  [[nodiscard]] const layer_info& get_info(int layer_id) const { return *info_[layer_id]; }

  /** stream used for allocations */
  [[nodiscard]] const cuda::stream& get_stream() const { return stream_; }

  /** max batch size (parameter given to constructor) */
  [[nodiscard]] IdxT max_batch_size() const { return batch_size_; }

  /** whether this node batcher computes reverse graphs */
  [[nodiscard]] bool do_reverse() const { return do_reverse_; }

 protected:
  // please try to avoid looking too much at the absolute abomination of C++
  // happening here, it is unfortunately the only way to ensure that we can
  // access differently named fields of the Mgraph, using C++ 11 and GCC
  template <bool IS_CSR = USE_CSR, int DUMMY = 0>
  struct neighbor_field_manager {};
  template <int DUMMY>
  struct neighbor_field_manager<false, DUMMY> {
    using nb_t = node_batcher<IdxT, false>;
    explicit neighbor_field_manager(nb_t& nb) : nb_(nb) {}
    inline void alloc(utils::allocator& d_alloc, int idx)
    {
      auto& info  = *nb_.info_[idx];
      auto& graph = *info.graph;
      // neighborhoods
      graph.neighbors = reinterpret_cast<IdxT*>(
        d_alloc.alloc(sizeof(IdxT) * graph.n_out_nodes * graph.sample_size, nb_.stream_()));
      // neighbor counts
      graph.neighbor_counts =
        reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT) * graph.n_out_nodes, nb_.stream_()));
    }
    inline void dealloc(utils::allocator& d_alloc, int idx)
    {
      auto& graph = *nb_.info_[idx]->graph;
      d_alloc.dealloc(graph.neighbors, 0, nb_.stream_());
      d_alloc.dealloc(graph.neighbor_counts, 0, nb_.stream_());
    }

   protected:
    nb_t& nb_;
  };
  template <int DUMMY>
  struct neighbor_field_manager<true, DUMMY> {
    using nb_t = node_batcher<IdxT, true>;
    explicit neighbor_field_manager(nb_t& nb) : nb_(nb) {}
    inline void alloc(utils::allocator& d_alloc, int idx)
    {
      auto& info  = *nb_.info_[idx];
      auto& graph = *info.graph;
      // CSR offsets
      graph.offsets = reinterpret_cast<IdxT*>(
        d_alloc.alloc(sizeof(IdxT) * (graph.n_out_nodes + 1), nb_.stream_()));
      // CSR indices
      graph.indices = reinterpret_cast<IdxT*>(
        d_alloc.alloc(sizeof(IdxT) * graph.n_out_nodes * graph.sample_size, nb_.stream_()));
    }
    inline void dealloc(utils::allocator& d_alloc, int idx)
    {
      auto& graph = *nb_.info_[idx]->graph;
      d_alloc.dealloc(graph.indices, 0, nb_.stream_());
      d_alloc.dealloc(graph.offsets, 0, nb_.stream_());
    }

   protected:
    nb_t& nb_;
  };

  inline void dealloc_layer_info_fields(utils::allocator& d_alloc,
                                        utils::allocator& h_alloc,
                                        int idx)
  {
    neighbor_field_manager<>(*this).dealloc(d_alloc, idx);
    // for all other indices, in_nodes is de-allocated through the MFG out_nodes
    if (idx == 0) d_alloc.dealloc(info_[0]->in_nodes, 0, stream_());
    d_alloc.dealloc(info_[idx]->graph->out_nodes, 0, stream_());
    h_alloc.dealloc(info_[idx]->n_in_nodes, 0, stream_());
    if (do_reverse_) {
      // reverse graph components
      d_alloc.dealloc(info_[idx]->graph_rev->indices, 0, stream_());
      d_alloc.dealloc(info_[idx]->graph_rev->offsets, 0, stream_());
      d_alloc.dealloc(info_[idx]->graph_rev->in_node_ids, 0, stream_());
    }
  }

  // this is called before renumbering, allowing sub-classes to copy
  // additional information from the original graph into the MFGs
  // by default, we don't do anything here
  virtual void copy_additional_info() {}

  // renumber the neighborhood and output node id's as per the node id's found
  // in the input to each layer.
  //
  // This way, the nodes at the first layer input represent a first level
  // "global" mapping for all subsequent layers
  //
  // This is done this way in order to keep the hidden embeddings to be able
  // to be stored in consecutive fashion for each layer.
  void renumber_node_ids(IdxT /*batch_size*/)
  {
    auto last_id = static_cast<int>(info_.size() - 1);
    for (auto i = last_id; i >= 0; --i) {
      auto* out_ptr = i == last_id ? batch_ : info_[i]->graph->out_nodes;
      auto* counts  = do_reverse_ ? info_[i]->graph_rev->offsets : nullptr;
      graph::renumber(*info_[i]->graph, counts, info_[i]->in_nodes, out_ptr, stream_());
      // at this point, all values for the MFG should be set
      if (do_reverse_) {
        graph::get_reverse_graph(*info_[i]->graph_rev,
                                 *info_[i]->graph,
                                 info_[i]->n_indices,
                                 workspace_,
                                 workspace_size_,
                                 stream_());
      }
    }
  }

  void assign_workspace()
  {
    auto& h_alloc = utils::host_allocator();
    auto& d_alloc = utils::device_allocator();
    auto last_id  = static_cast<int>(sample_sizes_.size() - 1);
    // for computing the output batch at each iteration
    next_start_ = reinterpret_cast<IdxT*>(h_alloc.alloc(sizeof(IdxT) * 2, stream_()));
    size_t wsize;
    IdxT* idx_null = nullptr;
    graph::get_next_output_batch(
      idx_null, idx_null, batch_size_, n_labels_, uint64_t{0}, nullptr, wsize, nullptr);
    // output batch (original id's)
    batch_    = reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT) * batch_size_, stream_()));
    IdxT size = batch_size_;
    // output batch (renumbered id's)
    info_[last_id]->graph->out_nodes =
      reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT) * batch_size_, stream_()));
    // each layer's storage
    for (auto i = last_id; i >= 0; --i) {
      auto& graph    = info_[i]->graph;
      auto prev_size = size;
      // max possible number of output nodes in this layer
      info_[i]->max_out_nodes = prev_size;
      size *= (sample_sizes_[i] + 1);  // +1 for self-loop
      auto n_size = info_[i]->max_out_nodes * sample_sizes_[i];
      // output nodes (same as input nodes to the layer above)
      if (i != last_id) { graph->out_nodes = info_[i + 1]->in_nodes; }
      // sampling factor
      info_[i]->sample_size = sample_sizes_[i];
      graph->sample_size    = sample_sizes_[i];
      // for now, set n_out_nodes to upper bound
      graph->n_out_nodes = info_[i]->max_out_nodes;
      // TODO(mjoux) at some point, group these allocations but ensure good alignments
      neighbor_field_manager<>(*this).alloc(d_alloc, i);
      // host allocations for calculating:
      // - number of unique nodes in each layer's input
      // - number of edges (nnz) if it's CSR or we have a reverse graph
      info_[i]->n_in_nodes = reinterpret_cast<IdxT*>(h_alloc.alloc(sizeof(IdxT) * 2, stream_()));
      info_[i]->n_indices  = info_[i]->n_in_nodes + 1;
      // uniquified node-id's
      info_[i]->in_nodes = reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT) * size, stream_()));
      size_t wsize_sample, wsize_uniq;
      graph::uniform_sample(*graph,
                            rng_,
                            graph_,
                            nullptr,
                            sample_type_,
                            IdxT{0} /*max degree unused here*/,
                            sample_self_,
                            info_[i]->n_indices,
                            nullptr,
                            wsize_sample,
                            nullptr);
      wsize = std::max(wsize_sample, wsize);
      graph::uniquify(nullptr,
                      info_[i]->n_in_nodes,
                      *graph,
                      nullptr,
                      prev_size,
                      n_ + 1,
                      nullptr,
                      wsize_uniq,
                      nullptr);
      wsize = std::max(wsize_uniq, wsize);
      if (do_reverse_) {
        auto& graph_rev = info_[i]->graph_rev;
        // set values for `graph_rev`
        graph_rev->n_out_nodes = info_[i]->max_out_nodes;
        graph_rev->sample_size = sample_sizes_[i];
        size_t rev_wsize_i;
        graph::get_reverse_graph(
          *graph_rev, *graph, info_[i]->n_indices, nullptr, rev_wsize_i, stream_());
        wsize = std::max(rev_wsize_i, wsize);
        // reverse graph components
        graph_rev->in_node_ids =
          reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT) * size, stream_()));
        // +1 for extra element representing full sum
        graph_rev->offsets =
          reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT) * (size + 1), stream_()));
        graph_rev->indices =
          reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT) * n_size, stream_()));
      }
    }
    workspace_size_ = wsize;
    workspace_      = reinterpret_cast<char*>(d_alloc.alloc(sizeof(char) * wsize, stream_()));
    // compute max vertex degree and use it allocate fast_int_div parameters
    auto* max_deg = reinterpret_cast<IdxT*>(d_alloc.alloc(sizeof(IdxT), stream_()));
    graph::max_vertex_degree(max_deg, graph_, stream_());
    RAFT_CUDA_TRY(
      cudaMemcpyAsync(&max_degree_, max_deg, sizeof(IdxT), cudaMemcpyDeviceToHost, stream_()));
    stream_.sync();
    d_alloc.dealloc(max_deg, 0, stream_());
  }

  /** (sub)graph information for nodes in the current batch, across layers */
  std::vector<std::unique_ptr<layer_info>> info_;
  /** sample sizes for each layer */
  std::vector<IdxT> sample_sizes_;

  /** Host side random number generator (for xor_key) */
  std::mt19937_64 gen_;

  /** graph to be sub-sampled */
  csc_t& graph_;
  /** RAFT RngState object */
  raft::random::RngState& rng_;

  /**
   * contains indices of the elements that are part of current batch.
   * [on device] [len = `batch_size_`].
   * This buffer is NOT renumbered and thus contains the node id's from the
   * original graph, as-is.
   */
  IdxT* batch_{nullptr};
  /**
   * list of numbers read so far for from the bijection. [pinned] [len = `2`]
   */
  IdxT* next_start_{nullptr};
  /** shared workspace for `uniquify()` and `get_next_output_batch()` and
      `get_reverse_graph()` */
  char* workspace_{nullptr};
  size_t workspace_size_{0};

  /** cuda stream to schedule work on */
  cuda::stream stream_{nullptr};

  /** a 64b number for XOR operation */
  uint64_t xor_key_{uint64_t{0}};

  /** output layer batch size */
  IdxT batch_size_{0};
  /** total elements to be considered */
  IdxT n_{0};
  /** total number of labeled nodes (may be smaller than n_) */
  IdxT n_labels_{0};
  /** total number of batches (or iterations) */
  IdxT num_batches_{0};
  /** current batch count (eg: in an epoch) */
  IdxT count_{0};
  /** current batch size */
  IdxT curr_batch_size_{0};
  /** maximum vertex degree */
  IdxT max_degree_{0};

  /* whether or not to create a reverse graph. */
  bool do_reverse_{false};
  /* whether or not to sample explicit self loops. true by default */
  bool sample_self_{true};
  /** sampling type */
  graph::SamplingAlgoT sample_type_{graph::SamplingAlgoT::kReservoirAlgoR};
};  // struct node_batcher

/**
 * @brief Generates the batch of nodes for heterogeneous graphs
 *
 * @note See `node_batcher` for more information.
 *
 */
template <typename IdxT, bool USE_CSR = false>
struct node_batcher_hg : node_batcher<IdxT, USE_CSR> {
  using csc_hg_t = graph::csc_hg<IdxT>;
  using mfg_hg_t =
    std::conditional_t<USE_CSR, graph::mfg_csc_hg<IdxT>, graph::mfg_ellpack_hg<IdxT>>;
  using graph_rev_hg_t = graph::mfg_csc_rev_hg<IdxT>;
  using nb_t           = node_batcher<IdxT, USE_CSR>;
  using batch_info_t   = typename nb_t::batch_info;
  using layer_info_t   = typename nb_t::layer_info;

  /**
   * Node batcher for heterogeneous graphs
   *
   * @note sampling/reverse will copy over node and edge types for all MFGs.
   *       However, the input node types will not be copied over.
   *       These can be copied with a simple gather operation if necessary.
   *
   * TODO(mjoux) for now, edge information is not copied over in reverse!
   *       thus, reverse algos cannot access edge info
   *
   * @param[in]    _graph      the graph to be sub-sampled
   * @param[in]    _sizes      list of sample sizes at each layer of network.
   *                           Its length will tell the number of layers in the
   *                           network. Sample sizes must be mentioned starting
   *                           from the first layer, all the way to the output
   *                           layer, in that order only.
   * @param[in]    _batch_size batch size
   * @param[in]    _n          total number of nodes
   * @param[in]    _n_labels   number of labeled nodes
   * @param[inout] _rng        RAFT RngState object
   * @param[in]    _type       sampling type
   * @param[in]    _do_reverse whether to compute reverse graph or not
   * @param[in]    _stream     cuda stream
   */
  node_batcher_hg(typename nb_t::csc_t& _graph,
                  const std::vector<IdxT>& _sizes,
                  IdxT _batch_size,
                  IdxT _n,
                  IdxT _n_labels,
                  raft::random::RngState& _rng,
                  graph::SamplingAlgoT _type,
                  bool _do_reverse,
                  const cuda::stream& _stream)
    : nb_t(_graph, _sizes, _batch_size, _n, _n_labels, _rng, _type, _do_reverse, _stream)
  {
    // overwrite the allocated MFGs from base class
    for (size_t i = 0; i < nb_t::sample_sizes_.size(); ++i) {
      nb_t::info_[i]->graph = std::make_unique<mfg_hg_t>(*nb_t::info_[i]->graph);
      if (nb_t::do_reverse_)
        nb_t::info_[i]->graph_rev = std::make_unique<graph_rev_hg_t>(*nb_t::info_[i]->graph_rev);
    }
    // then fill in fields specific to the HG MFG types
    assign_workspace_hg();
  }

  ~node_batcher_hg()
  {
    auto& h_alloc = utils::host_allocator();
    auto& d_alloc = utils::device_allocator();
    for (size_t i = 0; i < nb_t::info_.size(); ++i) {
      dealloc_layer_info_fields_hg(d_alloc, h_alloc, i);
    }
  }

  /** default copy constructor and copy assignment operator */
  node_batcher_hg(const node_batcher_hg&)            = default;
  node_batcher_hg& operator=(const node_batcher_hg&) = default;

 protected:
  inline mfg_hg_t* get_mfg_err(int idx)
  {
    auto* graph = dynamic_cast<mfg_hg_t*>(nb_t::info_[idx]->graph.get());
    if (graph == nullptr) CUGRAPH_OPS_ERROR("node_batcher_hg: MFG not a heterogeneous MFG");
    return graph;
  }

  inline mfg_hg_t* get_mfg(int idx)
  {
    auto* graph = dynamic_cast<mfg_hg_t*>(nb_t::info_[idx]->graph.get());
    ASSERT(graph != nullptr, "node_batcher_hg: MFG not a heterogeneous MFG");
    return graph;
  }

  inline graph_rev_hg_t* get_graph_rev_err(int idx)
  {
    auto* graph_rev = dynamic_cast<graph_rev_hg_t*>(nb_t::info_[idx]->graph_rev.get());
    if (graph_rev == nullptr)
      CUGRAPH_OPS_ERROR("node_batcher_hg: reverse MFG not a heterogeneous MFG");
    return graph_rev;
  }

  inline graph_rev_hg_t* get_graph_rev(int idx)
  {
    auto* graph_rev = dynamic_cast<graph_rev_hg_t*>(nb_t::info_[idx]->graph_rev.get());
    ASSERT(graph_rev != nullptr, "node_batcher_hg: reverse MFG not a heterogeneous MFG");
    return graph_rev;
  }

  void dealloc_layer_info_fields_hg(utils::allocator& d_alloc,
                                    utils::allocator& /*h_alloc*/,
                                    int idx)
  {
    auto* graph = get_mfg_err(idx);
    if (graph == nullptr) return;
    if (idx == 0) d_alloc.dealloc(graph->in_node_types, 0, nb_t::stream_());
    d_alloc.dealloc(graph->out_node_types, 0, nb_t::stream_());
    d_alloc.dealloc(graph->edge_types, 0, nb_t::stream_());
    if (nb_t::do_reverse_) {
      auto* graph_rev = get_graph_rev_err(idx);
      if (graph_rev == nullptr) return;
      d_alloc.dealloc(graph_rev->edge_types, 0, nb_t::stream_());
    }
  }

  void assign_workspace_hg()
  {
    auto& d_alloc = utils::device_allocator();
    auto last_id  = static_cast<int>(nb_t::sample_sizes_.size() - 1);
    for (int i = last_id; i >= 0; --i) {
      auto* graph = get_mfg(i);
      // set n_node_types/n_edge_types
      auto* full_graph = dynamic_cast<csc_hg_t*>(std::addressof(nb_t::graph_));
      ASSERT(full_graph != nullptr, "node_batcher_hg: full graph must be a heterogeneous graph");
      graph->n_edge_types = full_graph->n_edge_types;
      graph->n_node_types = full_graph->n_node_types;
      // n_out_nodes should be set to max_out_nodes by node batcher
      auto n_max_edges = graph->n_out_nodes * nb_t::sample_sizes_[i];
      graph->edge_types =
        reinterpret_cast<int32_t*>(d_alloc.alloc(sizeof(int32_t) * n_max_edges, nb_t::stream_()));
      if (i == last_id) {
        graph->out_node_types = reinterpret_cast<int32_t*>(
          d_alloc.alloc(sizeof(int32_t) * graph->n_out_nodes, nb_t::stream_()));
      } else {
        auto* graph_next      = get_mfg(i + 1);
        graph->out_node_types = graph_next->in_node_types;
      }
      auto n_in_nodes = graph->n_out_nodes * (nb_t::sample_sizes_[i] + 1);
      graph->in_node_types =
        reinterpret_cast<int32_t*>(d_alloc.alloc(sizeof(int32_t) * n_in_nodes, nb_t::stream_()));
      if (nb_t::do_reverse_) {
        auto* graph_rev         = get_graph_rev(i);
        graph_rev->n_edge_types = graph->n_edge_types;
        graph_rev->n_node_types = graph->n_node_types;

        graph_rev->edge_types =
          reinterpret_cast<int32_t*>(d_alloc.alloc(sizeof(int32_t) * n_max_edges, nb_t::stream_()));
        graph_rev->out_node_types = graph->out_node_types;
        graph_rev->in_node_types  = graph->in_node_types;
      }
    }
  }
};

}  // namespace cugraph::ops
