/*
 * Copyright (c) 2021-2022, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <raft/error.hpp>

#include <cstdint>

namespace cugraph::ops::gnn {
namespace graph {

/**
 * @brief constant to represent invalid node id.
 *
 * @tparam IdxT node index type
 *
 * @note It is important that its bit representation has 1 set everywhere,
 *       no matter which index type is used
 * @note As of now, we expect `IdxT` to be a signed integer type, although
 *       everything should work with unsigned types. In any case, we restrict
 *       the range of valid IDs to [0, 2^(sizeof(IdxT) * 8 - 1) - 2].
 *       This reduced range is due to reserving the most significant bit
 *       for (potential) indicator variables in algos, as well as one
 *       additional value for the invalid ID. This comes naturally when
 *       using signed integer types since the range simply excludes all
 *       negative values as well as the maximal positive value.
 */
template <typename IdxT>
static constexpr IdxT INVALID_ID = IdxT(-1);

/**
 * @brief constant to represent default node type
 *
 * @note We expect node and edge types to be representable with signed 32-bit
 *       integers, meaning that the valid range is [0, 2^31 - 1].
 */
static constexpr int32_t DEFAULT_NODE_TYPE = 0;

/**
 * @brief constant to represent default edge type
 *
 * @note Valid range: [0, 2^31 - 1] (see `DEFAULT_NODE_TYPE`)
 */
static constexpr int32_t DEFAULT_EDGE_TYPE = 0;

/**
 * @brief CSR representation of a full-graph.
 *
 * @tparam IdxT node index type (signed type, see `INVALID_ID`)
 *
 * @note By default, we expect an in-graph for most operations: an index into
 *       the `offsets` represents the node ID of a destination node and the
 *       values of `indices` are node IDs of source nodes. Operations requiring
 *       a "reverse" graph require the out-graph additionally.
 * @note Valid IDs: [0, 2^(sizeof(IdxT) * 8 - 1) - 2] (see `INVALID_ID`)
 * @note This object does NOT own any of the underlying pointers and thus is
 *       left to the calling code for their lifetime management.
 */
template <typename IdxT>
struct fg_csr {
  using idx_t = IdxT;
  // handle sub-classes correctly
  fg_csr()          = default;
  virtual ~fg_csr() = default;

  /** number of nodes belonging to this graph */
  IdxT n_nodes{0};
  /** total number of edges in this graph (length of `indices` array) */
  IdxT n_indices{0};
  /**
   * monotonically increasing array with each location pointing to the start
   * offset of the neighborhood of that node in the `indices` array.
   * It is of length `n_nodes + 1`.
   */
  IdxT* offsets{nullptr};
  /**
   * contains neighbor indices of every node belonging to this object.
   * It is of length `n_indices`.
   */
  IdxT* indices{nullptr};
  /**
   * edge feature indices for each of the above edge, pointing to that row in the
   * edge embedding table. This can be used to avoid having to reorder edge features
   * when we sort edges while creating CSR, or apply other transformations on
   * the graph structure.
   * Its length is `n_indices`.
   * @note This pointer is optional for many operations. If not given, it is
   *       assumed that edge features are ordered the same way as `indices`.
   * @note If provided, we assume that this is a perfect partition of the
   *       indexes in [0, n_indices - 1]. Operations using edge features only
   *       as input can usually work with any index, but operations writing
   *       to the edge features assume this is a partition unless specified.
   *       Otherwise, be aware that this could lead to data races!
   */
  IdxT* ef_indices{nullptr};
  /**
   * For every edge, this provides the indices of the reverse edges. Its length is
   * `n_indices`.
   *
   * @note For undirected graphs, these indices should point into the `indices`
   *       array of this same instance. For directed graphs, we expect two
   *       two instances of `fg_csr`, one for each direction, and these indices
   *       point into the `indices` array of the other instance/direction.
   * @note If one needs to know the edge embedding table row index of the reverse
   *       edge, then it can be computed as: `ef_indices[rev_edge_ids[curr_edge_id]]`
   * @note This pointer is optional for many operations, which only require
   *       the in-graph. It is not optional for operations requiring the out-graph.
   */
  IdxT* rev_edge_ids{nullptr};
};

/**
 * @brief CSR representation of a sequence of full-graphs.
 *
 * IOW, this is basically a 2-level indirection of arrays, as opposed to a
 * 1-level indirection in the case of classical CSR arrays. `graph_offsets`
 * array points to the start of locations of the `node_offsets` array for
 * every graph in this batch. In turn, `node_offsets` points to the neighborhood
 * information of each of the node in each of those graphs in the batch.
 *
 * However, the node id's across all graphs are given monotonically increasing
 * unique id's. That way, could, still use the `offsets` and `indices` arrays in
 * this class and use them as one big graph.
 *
 * This means that all indexes and values of the fg_csr base class (`offsets`,
 * `indices`, `ef_indices` and `rev_edge_ids` arrays) remain global IDs.
 * For the `ef_indices`, it is assumed to point to a global embedding table,
 * containing edge features of all edges in all graphs.
 *
 * The following snippet shows how one would access each graph's properties,
 * given the index of the graph.
 * @code{.cpp}
 * <init CSR sequence "graph_seq"...>
 * auto start_nodes = graph_seq.graph_offsets[graph_id];
 * auto end_nodes = graph_seq.graph_offsets[graph_id + 1];
 * auto n_nodes_in_graph = end_nodes - start_nodes;
 * auto n_edges_in_graph = graph_seq.offsets[end_nodes] - graph_seq.offsets[start_nodes];
 * @endcode
 *
 * @tparam IdxT node index type
 *
 * @note This object does NOT own any of the underlying pointers and thus is
 *       left to the calling code for their lifetime management.
 * @note By default, we expect the graphs to be in-graphs. See `fg_csr`.
 */
template <typename IdxT>
struct fg_csr_seq : public fg_csr<IdxT> {
  // handle sub-classes correctly
  fg_csr_seq() = default;

  /** number of graphs being part of this batch */
  IdxT n_graphs{0};
  /**
   * offsets representing start of the CSR info for each graph in the
   * `offsets` array. It is a monotonically increasing array. Its length is
   * `n_graphs + 1`
   */
  IdxT* graph_offsets{nullptr};
};  // struct fg_csr_seq

/**
 * @brief CSR representation of an individual batch of full graphs.
 *
 * This structure contains information to allow indexing of batch-local node
 * and edge features, given a `fg_csr_seq`.
 *
 * The following snippet shows how to obtain batch-local indexes, given
 * the index of a graph within the batch.
 * @code{.cpp}
 * <init graph batcher "batcher"...>
 * auto batch = batcher.next_batch();
 * <pick a batch idx "idx">
 * auto start_nodes = batch.csr_seq.graph_offsets[batch_offsets[idx]];
 * auto end_nodes = batch.csr_seq.graph_offsets[batch_offsets[idx] + 1];
 * auto last_node_global = end_nodes - 1;
 * auto last_node_local = last_node_global - start_nodes;
 * auto last_node_batch = last_node_global + batch.nf_offsets[idx];
 * auto last_edge_global = batch.csr_seq.offsets[end_nodes] - 1;
 * auto last_edge_local = last_edge_global - batch.csr_seq.offsets[end_nodes];
 * auto last_edge_batch = last_edge_global + batch.ef_offsets[idx];
 * @endcode
 *
 * @tparam IdxT node index type
 *
 * @note This object does NOT own any of the underlying pointers and thus is
 *       left to the calling code for their lifetime management.
 * @note By default, we expect the graphs to be in-graphs. See `fg_csr`.
 */
template <typename IdxT>
struct fg_csr_batch {
  // handle sub-classes correctly
  fg_csr_batch() = delete;

  explicit fg_csr_batch(const fg_csr_seq<IdxT>& csr_seq_) : csr_seq(csr_seq_) {}

  ~fg_csr_batch() = default;

  /** reference to the global `fg_csr_seq` structure */
  const fg_csr_seq<IdxT>& csr_seq;
  /** number of graphs part of this (individual) batch : the batch size */
  IdxT batch_size{0};
  /** total number of nodes in this (individual) batch */
  IdxT n_nodes{0};
  /** total number of edges in this (individual) batch */
  IdxT n_edges{0};
  /**
   * offsets within `csr_seq.graph_offsets` of the chosen graphs.
   * This is of length `batch_size`.
   */
  IdxT* batch_offsets{nullptr};
  /**
   * offsets of node features within a compacted (dense) batch-local
   * node feature matrix.
   * This is used to transform the (global) indices of `csr_seq.offsets`
   * into batch-local offsets.
   * This is of length `batch_size`.
   */
  IdxT* nf_offsets{nullptr};
  /**
   * offsets of edge features within a compacted (dense) batch-local
   * edge feature matrix.
   * This is used to transform the (global) indices of `csr_seq.indices`
   * into batch-local offsets.
   * This is of length `batch_size`.
   */
  IdxT* ef_offsets{nullptr};
};  // struct fg_csr_batch

/**
 * @brief ELLPACK representation of a MFG (message flow graph).
 *
 * @tparam IdxT node index type
 *
 * See DGL's definition of a MFG:
 * https://docs.dgl.ai/en/0.7.x/tutorials/large/L0_neighbor_sampling_overview.html
 *
 * We use MFGs here because they are bipartite and thus have a clear definition
 * of input/output nodes, which simplifies readability of code.
 *
 * ELLPACK representation means that we have a dense CSR-like adjacency matrix,
 * where the rows represent the output nodes of the MFG, and columns represent
 * the neighboring input nodes for a given output node.
 * Invalid node IDs indicate whether an output node has fewer neighbors than
 * the number of columns, but they must be placed at the end of the row!
 *
 * @note The input nodes consist of the union of output nodes and all their
 *       neighbors. Because of this, some input nodes may not appear in the
 *       adjacency matrix (`neighbors`).
 * @note Valid IDs: [0, 2^(sizeof(IdxT) * 8 - 1) - 2] (see `INVALID_ID`)
 * @note This object does NOT own any of the underlying pointers and thus is
 *       left to the calling code for their lifetime management.
 * @note This is always an in-graph (`neighbors` over the destination nodes).
 *       See `fg_csr` for further explanation on this.
 */
template <typename IdxT>
struct mfg_ellpack {
  // handle sub-classes correctly
  mfg_ellpack()          = default;
  virtual ~mfg_ellpack() = default;

  /** number of output nodes of this MFG. */
  IdxT n_out_nodes{0};
  /**
   * sample size: number of columns in `neighbors`.
   * This is the maximum number of neighbors of an output node
   */
  IdxT sample_size{0};
  /**
   * number of input nodes of this MFG.
   * An upper bound for this is `n_out_nodes * (sample_size + 1)`
   */
  IdxT n_in_nodes{0};
  /**
   * Output nodes of this MFG. [len = `n_out_nodes`]
   * For each output node, contains the index w.r.t. input nodes/features.
   */
  IdxT* out_nodes{nullptr};
  /**
   * Input nodes of this MFG. [len = `n_in_nodes`]
   * For each input node, contains the index w.r.t. input nodes/features of
   * the previous layer or the global graph if no previous layer exists.
   */
  IdxT* in_nodes{nullptr};
  /**
   * ELLPACK CSR-like adjacency matrix. [`n_out_nodes x sample_size`, device]
   * For every output node, this represents the neighbor IDs that were
   * sampled in the sampling process. It excludes implicit self-loops.
   */
  IdxT* neighbors{nullptr};
  /**
   * Neighborhood counts. [`n_out_nodes`, device]
   * For every output node, the number of neighbors
   * (excluding implicit self-loops).
   * If `neighbor_counts[n] == x`, the the `n`-th row of `neighbors` must
   * contain all valid neighbors of node `n` in the first `x` entries.
   */
  IdxT* neighbor_counts{nullptr};
  /** This is not supported yet. If `true`, `neighbors` is sorted by rows. */
  bool is_sorted{false};
};

/**
 * @brief CSR representation of a MFG.
 *
 * @tparam IdxT node index type
 *
 * Same as `mfg_ellpack` except the neighbors are in CSR format.
 *
 * @note Valid IDs: [0, 2^(sizeof(IdxT) * 8 - 1) - 2] (see `INVALID_ID`)
 * @note This object does NOT own any of the underlying pointers and thus is
 *       left to the calling code for their lifetime management.
 * @note This is always an in-graph (`offsets` over the destination nodes).
 *       See `fg_csr` for further explanation on this.
 */
template <typename IdxT>
struct mfg_csr {
  // handle sub-classes correctly
  mfg_csr()          = default;
  virtual ~mfg_csr() = default;

  /**
   * number of output nodes of this MFG. See `mfg_ellpack`.
   * Note that input/output nodes are defined as in the regular MFG
   */
  IdxT n_out_nodes{0};
  /** sample size: See `mfg_ellpack` */
  IdxT sample_size{0};
  /**
   * number of input nodes of this MFG. See `mfg_ellpack`
   * Note that input/output nodes are defined as in the regular MFG
   */
  IdxT n_in_nodes{0};
  /**
   * total number of edges in this graph (length of `indices` array)
   */
  IdxT n_indices{0};
  /**
   * Output nodes of this MFG. [len = `n_out_nodes`]
   * For each output node, contains the index w.r.t. input nodes/features.
   */
  IdxT* out_nodes{nullptr};
  /**
   * Input nodes of this MFG. [len = `n_in_nodes`]
   * For each input node, contains the index w.r.t. input nodes/features of
   * the previous layer or the global graph if no previous layer exists.
   */
  IdxT* in_nodes{nullptr};
  /**
   * monotonically increasing array with each location pointing to the start
   * offset of the neighborhood of that node in the `indices` array. It is of
   * length `n_out_nodes + 1`.
   */
  IdxT* offsets{nullptr};
  /**
   * contains neighbor indices of every node belonging to this object. It is of
   * length `n_indices`. For each node, the number of neighbors is
   * `<= sample_size`
   */
  IdxT* indices{nullptr};
};

/**
 * @brief CSR representation of a reversed/transposed MFG.
 *
 * @tparam IdxT node index type
 *
 * The reason why it is useful to distinguish this representation from
 * the full-graph csr is because an MFG is a bipartite graph, and we refer
 * to input/output nodes as in a normal MFG.
 * This means that what we call input nodes here are the nodes with ingoing
 * edges, and the output nodes have outgoing edges, since the graph is reversed.
 *
 * @note Because some of the input nodes of an MFG only appear due to keeping
 *       a reference to the output nodes, some of the nodes in this format
 *       may not have any neighbors!
 * @note Valid IDs: [0, 2^(sizeof(IdxT) * 8 - 1) - 2] (see `INVALID_ID`)
 * @note This object does NOT own any of the underlying pointers and thus is
 *       left to the calling code for their lifetime management.
 * @note This is always an out-graph (`offsets` over the source nodes).
 *       See `fg_csr` for further explanation on this.
 */
template <typename IdxT>
struct mfg_csr_rev {
  // handle sub-classes correctly
  mfg_csr_rev()          = default;
  virtual ~mfg_csr_rev() = default;

  /**
   * number of output nodes of this MFG. See `mfg_ellpack`.
   * Note that input/output nodes are defined as in the regular MFG
   */
  IdxT n_out_nodes{0};
  /** sample size: See `mfg_ellpack` */
  IdxT sample_size{0};
  /**
   * number of input nodes of this MFG. See `mfg_ellpack`
   * Note that input/output nodes are defined as in the regular MFG
   */
  IdxT n_in_nodes{0};
  /**
   * total number of edges in this graph (length of `indices` array)
   */
  IdxT n_indices{0};
  /**
   * For each input node, the ID w.r.t. output nodes of the MFG.
   * This may be invalid if the input node does not appear as an output node.
   * It is of length `n_in_nodes`.
   */
  IdxT* in_node_ids{nullptr};
  /**
   * monotonically increasing array with each location pointing to the start
   * offset of the neighborhood of that node in the `indices` array. It is of
   * length `n_in_nodes + 1`.
   */
  IdxT* offsets{nullptr};
  /**
   * contains neighbor indices of every node belonging to this object. It is of
   * length `n_indices`.
   */
  IdxT* indices{nullptr};
};

/**
 * @brief COO representation of an MFG.
 *
 * @tparam IdxT node index type
 *
 * The COO representation on its own is fairly useless, but additionally,
 * this contains reference arrays to link the `mfg_ellpack` to a COO
 * representation, as well as (optionally) the `mfg_csr_rev`.
 *
 * @note Valid IDs: [0, 2^(sizeof(IdxT) * 8 - 1) - 2] (see `INVALID_ID`)
 * @note This object does NOT own any of the underlying pointers and thus is
 *       left to the calling code for their lifetime management.
 */
template <typename IdxT>
struct mfg_coo {
  // handle sub-classes correctly
  mfg_coo()          = default;
  virtual ~mfg_coo() = default;

  /** number of output nodes of this MFG. See `mfg_ellpack` */
  IdxT n_out_nodes{0};
  /** sample size: See `mfg_ellpack` */
  IdxT sample_size{0};
  /** number of input nodes of this MFG. See `mfg_ellpack` */
  IdxT n_in_nodes{0};
  /** total number of edges in this graph. See `mfg_csr_rev` */
  IdxT n_indices{0};
  /**
   * Edge matrix [`2 x n_indices`, device].
   * There are 2 rows: the first row represents the input nodes of the edges,
   * the second row the output nodes.
   */
  IdxT* edges{nullptr};
  /**
   * ELLPACK CSR-like adjacency matrix which follows the exact same format
   * as the `neighbors` array in `mfg_ellpack`. Its values are the column index
   * of the given edge w.r.t. `edges` matrix.
   */
  IdxT* neighbor_idx{nullptr};
  /**
   * CSR-like array which follows the exact same format as the `indices` array
   * in `mfg_csr_rev`. Its values are the column index of the given edge
   * w.r.t. `edges` matrix.
   */
  IdxT* rev_idx{nullptr};
};

/**
 * @brief CSR representation of a heterogeneous graph.
 *
 * @tparam IdxT node index type
 *
 * @note See `fg_csr` first. This struct adds information needed for
 *       heterogeneous graphs.
 *
 * @note Only integer node and edge types are supported as of now.
 * @note By default, we expect the graph to be an in-graph. See `fg_csr`.
 */
template <typename IdxT>
struct fg_csr_hg : fg_csr<IdxT> {
  using fg_t  = fg_csr<IdxT>;
  using idx_t = typename fg_csr<IdxT>::idx_t;

  // handle sub-classes correctly
  fg_csr_hg()           = default;
  ~fg_csr_hg() override = default;

  explicit fg_csr_hg(const fg_t& graph) : fg_t(graph) {}
  fg_csr_hg(const fg_t& graph,
            int32_t _n_node_types,
            int32_t _n_edge_types,
            int32_t* _node_types,
            int32_t* _edge_types)
    : fg_t(graph),
      n_node_types(_n_node_types),
      n_edge_types(_n_edge_types),
      node_types(_node_types),
      edge_types(_edge_types)
  {
  }

  /** number of node types in this graph */
  int32_t n_node_types{1};
  /** number of edge types in this graph */
  int32_t n_edge_types{1};
  /** node types of each node in the graph (length `n_nodes`) */
  int32_t* node_types{nullptr};
  /**
   * edge types of each edge in the graph. It is of length `n_indices` and is
   * expected to be ordered the same way as `indices`
   */
  int32_t* edge_types{nullptr};
};

/**
 * @brief CSR representation of a heterogeneous MFG (message flow graph).
 *
 * @tparam IdxT node index type
 *
 * @note See `mfg_csr` first. This struct adds information needed for
 *       heterogeneous graphs.
 *
 * @note Only integer node and edge types are supported as of now.
 * @note This is always an in-graph, see `mfg_csr`.
 */
template <typename IdxT>
struct mfg_csr_hg : mfg_csr<IdxT> {
  using mfg_t = mfg_csr<IdxT>;

  // handle sub-classes correctly
  mfg_csr_hg()           = default;
  ~mfg_csr_hg() override = default;

  explicit mfg_csr_hg(const mfg_t& mfg) : mfg_t(mfg) {}
  mfg_csr_hg(const mfg_t& mfg,
             int32_t _n_node_types,
             int32_t _n_edge_types,
             int32_t* _out_node_types,
             int32_t* _in_node_types,
             int32_t* _edge_types)
    : mfg_t(mfg),
      n_node_types(_n_node_types),
      n_edge_types(_n_edge_types),
      out_node_types(_out_node_types),
      in_node_types(_in_node_types),
      edge_types(_edge_types)
  {
  }

  /** number of node types in this graph */
  int32_t n_node_types{1};
  /** number of edge types in this graph */
  int32_t n_edge_types{1};
  /**
   * node types of each output node of the MFG [`n_out_nodes`, device]
   * follows the same indexing as `out_nodes`
   */
  int32_t* out_node_types{nullptr};
  /**
   * node types of each input node of the MFG [`n_in_nodes`, device]
   * follows the same indexing as `in_nodes`
   */
  int32_t* in_node_types{nullptr};
  /**
   * edge types of each edge [`n_indices`, device]
   * follows the same indexing as `indices`
   */
  int32_t* edge_types{nullptr};
};

/**
 * @brief ELLPACK representation of a heterogeneous MFG (message flow graph).
 *
 * @tparam IdxT node index type
 *
 * @note See `mfg_ellpack` first. This struct adds information needed for
 *       heterogeneous graphs.
 *
 * @note Only integer node and edge types are supported as of now.
 * @note This is always an in-graph, see `mfg_ellpack`.
 */
template <typename IdxT>
struct mfg_ellpack_hg : mfg_ellpack<IdxT> {
  using mfg_t = mfg_ellpack<IdxT>;

  // handle sub-classes correctly
  mfg_ellpack_hg()           = default;
  ~mfg_ellpack_hg() override = default;

  explicit mfg_ellpack_hg(const mfg_t& mfg) : mfg_t(mfg) {}
  mfg_ellpack_hg(const mfg_t& mfg,
                 int32_t _n_node_types,
                 int32_t _n_edge_types,
                 int32_t* _out_node_types,
                 int32_t* _in_node_types,
                 int32_t* _edge_types)
    : mfg_t(mfg),
      n_node_types(_n_node_types),
      n_edge_types(_n_edge_types),
      out_node_types(_out_node_types),
      in_node_types(_in_node_types),
      edge_types(_edge_types)
  {
  }

  /** number of node types in this graph */
  int32_t n_node_types{1};
  /** number of edge types in this graph */
  int32_t n_edge_types{1};
  /**
   * node types of each output node of the MFG [`n_out_nodes`, device]
   * follows the same indexing as `out_nodes`
   */
  int32_t* out_node_types{nullptr};
  /**
   * node types of each input node of the MFG [`n_in_nodes`, device]
   * follows the same indexing as `in_nodes`
   */
  int32_t* in_node_types{nullptr};
  /**
   * edge types of each edge [`n_out_nodes x sample_size`, device]
   * follows the same indexing as `neighbors`
   */
  int32_t* edge_types{nullptr};
};

/**
 * @brief CSR representation of a reverse heterogeneous MFG.
 *
 * @tparam IdxT node index type
 *
 * @note See `mfg_csr_rev` first. This struct adds information needed for
 *       heterogeneous graphs.
 *
 * @note Only integer node and edge types are supported as of now.
 * @note This is always an out-graph, see `mfg_csr_rev`.
 */
template <typename IdxT>
struct mfg_csr_rev_hg : mfg_csr_rev<IdxT> {
  using mfg_t = mfg_csr_rev<IdxT>;

  // handle sub-classes correctly
  mfg_csr_rev_hg()           = default;
  ~mfg_csr_rev_hg() override = default;

  explicit mfg_csr_rev_hg(const mfg_t& mfg) : mfg_t(mfg) {}
  mfg_csr_rev_hg(const mfg_t& mfg,
                 int32_t _n_node_types,
                 int32_t _n_edge_types,
                 int32_t* _out_node_types,
                 int32_t* _in_node_types,
                 int32_t* _edge_types)
    : mfg_t(mfg),
      n_node_types(_n_node_types),
      n_edge_types(_n_edge_types),
      out_node_types(_out_node_types),
      in_node_types(_in_node_types),
      edge_types(_edge_types)
  {
  }

  /** number of node types in this graph */
  int32_t n_node_types{1};
  /** number of edge types in this graph */
  int32_t n_edge_types{1};
  /**
   * node types of each output node of the MFG [`n_out_nodes`, device]
   * follows the same indexing as `out_nodes` in the non-reverse MFG
   */
  int32_t* out_node_types{nullptr};
  /**
   * node types of each input node of the MFG [`n_in_nodes`, device]
   * follows the same indexing as `in_node_ids` /
   * `in_nodes` in the non-reverse MFG
   */
  int32_t* in_node_types{nullptr};
  /**
   * edge types of each edge [`n_indices`, device]
   * follows the same indexing as `indices`
   */
  int32_t* edge_types{nullptr};
};

}  // namespace graph

/** Instantiations for different index types (in main namespace) */

/**
 * @brief Full-graph CSR representation for 32b node id's
 */
using fg_csr_s32_t = graph::fg_csr<int32_t>;
/**
 * @brief Full-graph CSR representation for 64b node id's
 */
using fg_csr_s64_t = graph::fg_csr<int64_t>;

/**
 * @brief Full-graph sequence of CSR representation for 32b node id's.
 */
using fg_csr_seq_s32_t = graph::fg_csr_seq<int32_t>;

/**
 * @brief Full-graph sequence of CSR representation for 64b node id's
 */
using fg_csr_seq_s64_t = graph::fg_csr_seq<int64_t>;

/**
 * @brief Representation of a single batch of CSR full-graphs for 32b node id's.
 */
using fg_csr_batch_s32_t = graph::fg_csr_batch<int32_t>;

/**
 * @brief Representation of a single batch of CSR full-graphs for 64b node id's.
 */
using fg_csr_batch_s64_t = graph::fg_csr_batch<int64_t>;

/**
 * @brief MFG ELLPACK representation for 32b node id's
 */
using mfg_ellpack_s32_t = graph::mfg_ellpack<int32_t>;
/**
 * @brief MFG ELLPACK representation for 64b node id's
 */
using mfg_ellpack_s64_t = graph::mfg_ellpack<int64_t>;

/**
 * @brief MFG CSR representation for 32b node id's
 */
using mfg_csr_s32_t = graph::mfg_csr<int32_t>;
/**
 * @brief MFG CSR representation for 64b node id's
 */
using mfg_csr_s64_t = graph::mfg_csr<int64_t>;

/**
 * @brief MFG reverse CSR representation for 32b node id's
 */
using mfg_csr_rev_s32_t = graph::mfg_csr_rev<int32_t>;
/**
 * @brief MFG reverse CSR representation for 64b node id's
 */
using mfg_csr_rev_s64_t = graph::mfg_csr_rev<int64_t>;

/**
 * @brief MFG COO representation for 32b node id's
 */
using mfg_coo_s32_t = graph::mfg_coo<int32_t>;
/**
 * @brief MFG COO representation for 64b node id's
 */
using mfg_coo_s64_t = graph::mfg_coo<int64_t>;

/**
 * @brief Full-graph CSR representation of a heterogeneous graph (32 id's)
 */
using fg_csr_hg_s32_t = graph::fg_csr_hg<int32_t>;
/**
 * @brief Full-graph CSR representation of a heterogeneous graph (64b id's)
 */
using fg_csr_hg_s64_t = graph::fg_csr_hg<int64_t>;

/**
 * @brief MFG ELLPACK representation of a heterogeneous graph (32 id's)
 */
using mfg_ellpack_hg_s32_t = graph::mfg_ellpack_hg<int32_t>;
/**
 * @brief MFG ELLPACK representation of a heterogeneous graph (64b id's)
 */
using mfg_ellpack_hg_s64_t = graph::mfg_ellpack_hg<int64_t>;

/**
 * @brief MFG CSR representation of a heterogeneous graph (32 id's)
 */
using mfg_csr_hg_s32_t = graph::mfg_csr_hg<int32_t>;
/**
 * @brief MFG CSR representation of a heterogeneous graph (64b id's)
 */
using mfg_csr_hg_s64_t = graph::mfg_csr_hg<int64_t>;

/**
 * @brief MFG reverse CSR representation of a heterogeneous graph (32 id's)
 */
using mfg_csr_rev_hg_s32_t = graph::mfg_csr_rev_hg<int32_t>;
/**
 * @brief MFG reverse CSR representation of a heterogeneous graph (64b id's)
 */
using mfg_csr_rev_hg_s64_t = graph::mfg_csr_rev_hg<int64_t>;

}  // namespace cugraph::ops::gnn
