/*
 * 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 <catch2/catch.hpp>

#include <set>
#include <vector>

namespace cugraph::ops::catch2 {

/**
 * @brief Custom matcher class for comparing device arrays
 *        representing many multi-sets
 *
 * @tparam DataT data type of the arrays
 */
template <typename DataT, typename IdxT>
class device_multiset_equals_t : public Catch::MatcherBase<DataT*> {
  std::vector<DataT> h_ref_;
  std::vector<IdxT> h_idx_;

 public:
  device_multiset_equals_t(DataT* d_ref, size_t len, IdxT* d_idx, size_t n_idx)
  {
    h_ref_.resize(len);
    utils::copy(h_ref_.data(), d_ref, len, cudaMemcpyDeviceToHost);
    h_idx_.resize(n_idx);
    utils::copy(h_idx_.data(), d_idx, n_idx, cudaMemcpyDeviceToHost);
  }

  bool match(DataT* const& d_computed) const override
  {
    std::vector<DataT> h_computed(h_ref_.size());
    utils::copy(h_computed.data(), d_computed, h_ref_.size(), cudaMemcpyDeviceToHost);
    for (size_t i = 0; i < h_idx_.size() - 1; ++i) {
      IdxT set_off1 = h_idx_[i], set_off2 = h_idx_[i + 1];
      std::multiset<DataT> ref_set(h_ref_.begin() + set_off1, h_ref_.begin() + set_off2);
      std::multiset<DataT> computed_set(h_computed.begin() + set_off1,
                                        h_computed.begin() + set_off2);
      // multi-set != operator removed in C++ 20, but == operator should
      // be what we need here (test for size and same sorted elements)
      if (!(ref_set == computed_set)) {
        UNSCOPED_INFO("@i= " << i << ", set [" << set_off1 << ", " << set_off2
                             << "[ ref multi-set != val multi-set:");
        // length must be the same here since offset indexes are the same
        for (auto i = ref_set.begin(), j = computed_set.begin();
             i != ref_set.end() && j != computed_set.end();
             ++i, ++j) {
          UNSCOPED_INFO(*i << " | " << *j);
        }
        return false;
      }
    }
    return true;
  }

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

/**
 * @brief Builder function for the custom device array matcher
 *
 * @tparam DataT data type of the arrays
 *
 * @param[in] d_ref reference device array
 * @param[in] len   length of the array
 * @param[in] d_idx reference index array (ensure to test these indices
 *                  between reference and computed before using this matcher)
 * @param[in] n_idx length of the index array
 *
 * @return the matcher object
 */
template <typename DataT, typename IdxT>
device_multiset_equals_t<DataT, IdxT> device_multiset_equals(DataT* d_ref,
                                                             size_t len,
                                                             IdxT* d_idx,
                                                             size_t n_idx)
{
  return device_multiset_equals_t<DataT, IdxT>(d_ref, len, d_idx, n_idx);
}

}  // namespace cugraph::ops::catch2
