/*
 * 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 "test_utils.hpp"

#include <cugraph-ops/graph/format.hpp>
#include <cugraph-ops/graph/reverse_graph.hpp>
#include <cugraph-ops/test/random/csc_generator.hpp>
#include <cugraph-ops/test/test_utils.hpp>
#include <raft/random/rng_state.hpp>
#include <thrust/host_vector.h>

namespace cugraph::ops {

template <typename IdxT, bool IS_HG = false>
struct graph_buffers {
  device_buffer<IdxT> map_csc_to_coo;

  device_buffer<IdxT> rev_offsets;
  device_buffer<IdxT> rev_indices;
  device_buffer<IdxT> map_rev_to_coo;

  host_buffer<IdxT> h_offsets;
  host_buffer<IdxT> h_indices;
  host_buffer<IdxT> h_map_csc_to_coo;
  host_buffer<int32_t> h_edge_types;

  template <typename GraphT>
  void init_buffers_and_update_graph(GraphT& graph,
                                     raft::random::RngState& r,
                                     const int32_t* edge_types = nullptr)
  {
    if (h_offsets.length() > 0) {
      h_offsets.copy_from(graph.offsets, h_offsets.length());
      h_indices.copy_from(graph.indices, h_indices.length());

      if (IS_HG) { h_edge_types.copy_from(edge_types, h_edge_types.length()); }
    }

    // important: map_csc_to_coo before reversal, otherwise map_rev_to_coo is incorrect
    if (map_csc_to_coo.length() > 0) {
      device_buffer<IdxT> workspace;
      generate_perm_index(
        map_csc_to_coo.data(), workspace, r, static_cast<IdxT>(map_csc_to_coo.length()));
      graph.map_csc_to_coo = map_csc_to_coo;
      if (h_map_csc_to_coo.length() > 0) h_map_csc_to_coo.copy_from(map_csc_to_coo);
    }

    if (rev_offsets.length() > 0) {
      size_t cub_workspace_size = 0;
      graph.rev_offsets         = rev_offsets;
      graph.rev_indices         = rev_indices;
      graph.map_rev_to_coo      = map_rev_to_coo;
      device_buffer<IdxT> node_counts{graph.n_src_nodes};
      graph::get_reverse_graph(graph, node_counts, nullptr, cub_workspace_size, nullptr);
      device_buffer<char> cub_workspace{cub_workspace_size};
      graph::get_reverse_graph(graph, node_counts, cub_workspace, cub_workspace_size, nullptr);
    }
  }
};

template <typename IdxT, bool IS_HG = false>
graph_buffers<IdxT, IS_HG> create_graph_buffers(IdxT n_src_nodes,
                                                IdxT n_dst_nodes,
                                                IdxT n_indices,
                                                bool reverse_bwd,
                                                bool use_map_csc_to_coo,
                                                bool host)
{
  auto map_csc_to_coo_len = use_map_csc_to_coo ? n_indices : 0;
  auto rev_offsets_len    = reverse_bwd ? n_src_nodes + 1 : 0;
  auto rev_indices_len    = reverse_bwd ? n_indices : 0;

  auto h_offsets_len        = n_dst_nodes + 1;
  auto h_indices_len        = n_indices;
  auto h_map_csc_to_coo_len = use_map_csc_to_coo ? n_indices : 0;
  auto h_edge_types_len     = IS_HG ? n_indices : 0;

  return {// map
          device_buffer<IdxT>(map_csc_to_coo_len),
          // reverse
          device_buffer<IdxT>(rev_offsets_len),
          device_buffer<IdxT>(rev_indices_len),
          device_buffer<IdxT>(rev_indices_len),
          // host-side buffers
          host_buffer<IdxT>(host ? h_offsets_len : 0),
          host_buffer<IdxT>(host ? h_indices_len : 0),
          host_buffer<IdxT>(host ? h_map_csc_to_coo_len : 0),
          host_buffer<int32_t>(host ? h_edge_types_len : 0)};
}

template <typename IdxT>
void generate_csc(graph::csc<IdxT>& graph,
                  raft::random::RngState& r,
                  IdxT nodes,
                  IdxT nnz,
                  bool is_undirected = false)
{
  cuda::stream s(nullptr);
  generate_csc_on_device(graph, r, nodes, nnz, is_undirected, s);
}

template <typename IdxT>
void generate_csc(
  graph::csc<IdxT>& graph, IdxT* rev_edge_ids, raft::random::RngState& r, IdxT nodes, IdxT nnz)
{
  cuda::stream s(nullptr);
  generate_csc_on_device(graph, rev_edge_ids, r, nodes, nnz, s);
}

template <typename IdxT>
void delete_csc(const graph::csc<IdxT>& graph)
{
  cuda::stream s(nullptr);
  delete_csc_from_device(graph, s);
}

template <typename IdxT>
void generate_csc(graph::bipartite_csc<IdxT>& graph,
                  raft::random::RngState& r,
                  IdxT src_nodes,
                  IdxT dst_nodes,
                  IdxT nnz)
{
  cuda::stream s(nullptr);
  generate_csc_on_device(graph, r, src_nodes, dst_nodes, nnz, s);
}

template <typename IdxT>
void delete_csc(const graph::bipartite_csc<IdxT>& graph)
{
  cuda::stream s(nullptr);
  delete_csc_from_device(graph, s);
}

template <typename IdxT>
struct __attribute__((aligned(16))) ellpack_info {
  host_buffer<IdxT> neighbors{};
  host_buffer<IdxT> neighbor_counts{};
  ellpack_info() = default;
  ellpack_info(const ellpack_info& other)
  {
    neighbors.reserve(other.neighbors.size());
    neighbor_counts.reserve(other.neighbor_counts.size());
    neighbors.copy_from(other.neighbors);
    neighbor_counts.copy_from(other.neighbor_counts);
  }
};

template <typename IdxT>
struct __attribute__((aligned(16))) csc_info {
  host_buffer<IdxT> offsets{};
  host_buffer<IdxT> indices{};
  csc_info() = default;
  csc_info(const csc_info& other)
  {
    offsets.reserve(other.offsets.size());
    indices.reserve(other.indices.size());
    offsets.copy_from(other.offsets);
    indices.copy_from(other.indices);
  }
};

}  // namespace cugraph::ops
