/*
 * Copyright (c) 2022-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 <algorithm>
#include <vector>

namespace cugraph::ops::catch2 {

/**
 * @brief Custom matcher class for matching sampled indexes
 *
 * @tparam DataT data type of the arrays
 */
template <typename DataT>
class device_sampled_index_t : public Catch::MatcherBase<DataT*> {
  std::vector<DataT> h_sizes_;
  const int sample_size_;
  const bool replace_;

 public:
  explicit device_sampled_index_t(DataT* d_sizes, DataT n_sizes, int sample_size, bool replace)
    : sample_size_(sample_size), replace_(replace)
  {
    h_sizes_.resize(n_sizes);
    utils::copy(h_sizes_.data(), d_sizes, n_sizes, cudaMemcpyDeviceToHost);
  }

  bool match(DataT* const& d_indices) const override
  {
    if (sample_size_ <= 0) {
      UNSCOPED_INFO("Sample size <= 0 makes no sense for sampling.");
      return false;
    }
    auto sample_size = static_cast<size_t>(sample_size_);
    std::vector<DataT> h_indices(h_sizes_.size() * sample_size);
    utils::copy(h_indices.data(), d_indices, h_indices.size(), cudaMemcpyDeviceToHost);
    for (size_t i = 0; i < h_sizes_.size(); ++i) {
      size_t i_start = i * sample_size;
      auto size      = h_sizes_[i];
      for (size_t j = 0; j < sample_size; ++j) {
        auto bad = false;
        auto idx = h_indices[i_start + j];
        if (size <= 0 && idx != graph::INVALID_ID<DataT>) bad = true;
        if (replace_ && size > 0 && (idx < 0 || idx >= size)) bad = true;
        if (!replace_ && j >= size && idx != graph::INVALID_ID<DataT>) bad = true;
        if (bad) {
          UNSCOPED_INFO("@i=" << i << " j=" << j << " size=" << size << " idx=" << idx
                              << " (replace=" << replace_ << ")");
          return false;
        }
      }
      if (!replace_ && size > 0) {
        // check that the indexes are unique up to size are unique
        auto end   = static_cast<size_t>(std::min(size, DataT{sample_size_}));
        auto i_end = i * sample_size + end;
        auto uniq  = std::vector<DataT>(h_indices.begin() + i_start, h_indices.begin() + i_end);
        std::sort(uniq.begin(), uniq.end());
        for (size_t j = 1; j < end; ++j) {
          if (uniq[j] == uniq[j - 1]) {
            UNSCOPED_INFO("@i=" << i << " j=" << j << " duplicate (w/o replacement): "
                                << uniq[j - 1] << " == " << uniq[j]);
            return false;
          }
        }
      }
    }
    return true;
  }

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

/**
 * @brief Builder function for the custom device sampled indexes matcher
 *
 * @tparam DataT data type of the arrays
 *
 * @param[in]   d_sizes     sampled population sizes [device] [n_sizes]
 * @param[in]   n_sizes     number of population sizes
 * @param[in]   sample_size max number of samples per size
 * @param[in]   replace     whether sampling is done with replacement
 *
 * @return the matcher object
 */
template <typename DataT>
device_sampled_index_t<DataT> device_sampled_index(DataT* d_sizes,
                                                   DataT n_sizes,
                                                   int sample_size,
                                                   bool replace)
{
  return device_sampled_index_t<DataT>(d_sizes, n_sizes, sample_size, replace);
}

}  // namespace cugraph::ops::catch2
