/*
 * 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 <utils/cudart.hpp>

#include <catch2/catch.hpp>

#include <string>
#include <vector>

namespace cugraph::ops::catch2 {

/**
 * @brief Custom matcher class for checking sampling
 *
 * @tparam IdxT data type of the arrays
 */
template <typename IdxT>
class sample_matcher_t : public Catch::MatcherBase<IdxT*> {
  std::vector<IdxT> h_csc_offsets_;
  const IdxT sample_size_;
  std::vector<IdxT> h_sampled_nodes_;

 public:
  sample_matcher_t(
    IdxT* d_csc_offsets, size_t len, IdxT* d_sampled_nodes, size_t sampled_len, IdxT sample_size)
    : sample_size_(sample_size)
  {
    h_csc_offsets_.resize(len);
    utils::copy(h_csc_offsets_.data(), d_csc_offsets, len, cudaMemcpyDeviceToHost);
    h_sampled_nodes_.resize(sampled_len);
    if (d_sampled_nodes != nullptr) {
      utils::copy(h_sampled_nodes_.data(), d_sampled_nodes, sampled_len, cudaMemcpyDeviceToHost);
    } else {
      for (IdxT i = 0; i < sampled_len; ++i) {
        h_sampled_nodes_[i] = i;
      }
    }
  }

  bool match(IdxT* const& d_counts) const override
  {
    std::vector<IdxT> h_counts(h_sampled_nodes_.size());
    utils::copy(h_counts.data(), d_counts, h_sampled_nodes_.size(), cudaMemcpyDeviceToHost);
    for (size_t i = 0; i < h_sampled_nodes_.size(); ++i) {
      auto node_id = h_sampled_nodes_[i];
      auto count   = h_csc_offsets_[node_id + 1] - h_csc_offsets_[node_id];
      auto min     = std::min(count, sample_size_);
      if (h_counts[i] > min) return false;
    }
    return true;
  }

  [[nodiscard]] std::string describe() const override { return "sampling checker"; }
};  // class sample_matcher_t

/**
 * @brief Builder function for the custom sampling checker
 *
 * @tparam IdxT data type of the arrays
 *
 * @param[in] d_csc_offsets   CSR offsets for the nodes in the graph
 * @param[in] len             length of this array
 * @param[in] d_sampled_nodes sampled node counts
 * @param[in] sampled_len     length of this array
 * @param[in] sample_size     sample size
 *
 * @return the matcher object
 */
template <typename IdxT>
sample_matcher_t<IdxT> sample_matcher(
  IdxT* d_csc_offsets, size_t len, IdxT* d_sampled_nodes, size_t sampled_len, IdxT sample_size)
{
  return sample_matcher_t<IdxT>(d_csc_offsets, len, d_sampled_nodes, sampled_len, sample_size);
}

}  // namespace cugraph::ops::catch2
