/*
 * Copyright (c) 2021-2023, 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 <cugraph-ops/graph/format.hpp>

#include <cstdint>

namespace cugraph::ops {
namespace graph {

/**
 * @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 Mgraph, 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).
 *
 */
template <typename IdxT>
struct __attribute__((aligned(16))) mfg_ellpack {
  using idx_t = IdxT;
  // handle sub-classes correctly
  mfg_ellpack()          = default;
  virtual ~mfg_ellpack() = default;

  /**
   * 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};
  /**
   * 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};

  /** 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};

  /** 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 `csc` for further explanation on this.
 */
template <typename IdxT>
struct mfg_csc {
  using idx_t = IdxT;
  // handle sub-classes correctly
  mfg_csc()          = default;
  virtual ~mfg_csc() = 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};
  /**
   * 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 csc 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).
 *
 */
template <typename IdxT>
struct __attribute__((aligned(16))) mfg_csc_rev {
  using idx_t = IdxT;
  // handle sub-classes correctly
  mfg_csc_rev()          = default;
  virtual ~mfg_csc_rev() = default;

  /**
   * 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};

  /**
   * 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};
};

/**
 * @brief CSR representation of a heterogeneous MFG (message flow graph).
 *
 * @tparam IdxT node index type
 *
 * @note See `mfg_csc` 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_csc`.
 */
template <typename IdxT>
struct mfg_csc_hg : mfg_csc<IdxT> {
  using idx_t  = IdxT;
  using base_t = mfg_csc<IdxT>;

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

  explicit mfg_csc_hg(const base_t& graph) : base_t(graph) {}
  mfg_csc_hg(const base_t& graph,
             int32_t _n_node_types,
             int32_t _n_edge_types,
             int32_t* _out_node_types,
             int32_t* _in_node_types,
             int32_t* _edge_types)
    : base_t(graph),
      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)
  {
  }

  /**
   * 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]
   * indexed over the `n_in_nodes` range.
   */
  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};

  /** 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};
};

/**
 * @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 idx_t  = IdxT;
  using base_t = mfg_ellpack<IdxT>;

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

  explicit mfg_ellpack_hg(const base_t& mfg) : base_t(mfg) {}
  mfg_ellpack_hg(const base_t& graph,
                 int32_t _n_node_types,
                 int32_t _n_edge_types,
                 int32_t* _out_node_types,
                 int32_t* _in_node_types,
                 int32_t* _edge_types)
    : base_t(graph),
      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)
  {
  }

  /**
   * 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]
   * indexed over the `n_in_nodes` range.
   */
  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};

  /** 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};
};

/**
 * @brief CSR representation of a reverse heterogeneous MFG.
 *
 * @tparam IdxT node index type
 *
 * @note See `mfg_csc_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_csc_rev`.
 */
template <typename IdxT>
struct mfg_csc_rev_hg : mfg_csc_rev<IdxT> {
  using idx_t  = IdxT;
  using base_t = mfg_csc_rev<IdxT>;

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

  explicit mfg_csc_rev_hg(const base_t& graph) : base_t(graph) {}
  mfg_csc_rev_hg(const base_t& graph,
                 int32_t _n_node_types,
                 int32_t _n_edge_types,
                 int32_t* _out_node_types,
                 int32_t* _in_node_types,
                 int32_t* _edge_types)
    : base_t(graph),
      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)
  {
  }

  /**
   * 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]
   * indexed over the `n_in_nodes` range of 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};

  /** 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};
};

}  // namespace graph

/**
 * @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_csc_s32_t = graph::mfg_csc<int32_t>;
/**
 * @brief MFG CSR representation for 64b node id's
 */
using mfg_csc_s64_t = graph::mfg_csc<int64_t>;

/**
 * @brief MFG reverse CSR representation for 32b node id's
 */
using mfg_csc_rev_s32_t = graph::mfg_csc_rev<int32_t>;
/**
 * @brief MFG reverse CSR representation for 64b node id's
 */
using mfg_csc_rev_s64_t = graph::mfg_csc_rev<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_csc_hg_s32_t = graph::mfg_csc_hg<int32_t>;
/**
 * @brief MFG CSR representation of a heterogeneous graph (64b id's)
 */
using mfg_csc_hg_s64_t = graph::mfg_csc_hg<int64_t>;

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

}  // namespace cugraph::ops
