/*
 * Copyright (c) 2021-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 <utils/cudart.hpp>

#include <cugraph-ops/graph/format.hpp>

#include <catch2/catch.hpp>

#include <vector>

namespace cugraph::ops::catch2 {

/**
 * @brief Custom matcher class for comparing COO vs CSR src index arrays
 *
 * @tparam IdxT index type of the MFGs
 */
template <typename IdxT>
class device_coo_csc_equals_t : public Catch::MatcherBase<IdxT*> {
  std::vector<IdxT> h_ref_;
  std::vector<IdxT> h_ids_;
  IdxT n_idx_;

 public:
  explicit device_coo_csc_equals_t(IdxT* csc_offsets, IdxT* node_ids, IdxT n_dst_nodes)
  {
    h_ref_.resize(n_dst_nodes + 1);
    utils::copy(h_ref_.data(), csc_offsets, h_ref_.size(), cudaMemcpyDeviceToHost);
    h_ids_.resize(n_dst_nodes);
    utils::copy(h_ids_.data(), node_ids, h_ids_.size(), cudaMemcpyDeviceToHost);
    n_idx_ = h_ref_[n_dst_nodes];
  }

  bool match(IdxT* const& src_indices) const override
  {
    std::vector<IdxT> h_idx(n_idx_);
    utils::copy(h_idx.data(), src_indices, h_idx.size(), cudaMemcpyDeviceToHost);
    for (size_t i = 0; i < h_ids_.size(); ++i) {
      auto src_expected = h_ids_[i];
      for (IdxT j = h_ref_[i]; j < h_ref_[i + 1]; ++j) {
        auto src_computed = h_idx[j];
        if (src_expected != src_computed) {
          UNSCOPED_INFO("@node=" << i << " / idx=" << j << " expected ID " << src_expected
                                 << ", but got " << src_computed);
          return false;
        }
      }
    }
    return true;
  }

  [[nodiscard]] std::string describe() const override { return "device coo_csc matcher"; }
};  // class device_coo_csc_equals_t

/**
 * @brief Builder function for the custom COO vs CSR matcher
 *
 * @tparam IdxT index type of the MFGs
 *
 * @param[in] csc_offsets   reference CSR offsets
 * @param[in] node_ids      reference node IDs per CSR offset
 * @param[in] n_dst_nodes       number of nodes (size of `node_ids` and 1 less than
 *                          size of `csc_offsets`)
 *
 * @return the matcher object
 */
template <typename IdxT>
device_coo_csc_equals_t<IdxT> device_coo_csc_equals(IdxT* csc_offsets,
                                                    IdxT* node_ids,
                                                    IdxT n_dst_nodes)
{
  return device_coo_csc_equals_t<IdxT>(csc_offsets, node_ids, n_dst_nodes);
}

}  // namespace cugraph::ops::catch2
