/*
 * 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 <cmath>
#include <string>
#include <utility>
#include <vector>

namespace cugraph::ops::catch2 {

/**
 * @brief Custom matcher class for comparing device arrays
 *
 * @tparam DataT data type of the arrays
 */
template <typename DataT>
class device_equals_t : public Catch::MatcherBase<DataT*> {
  std::vector<DataT> h_ref_;

 public:
  device_equals_t(DataT* d_ref, size_t len)
  {
    h_ref_.resize(len);
    utils::copy(h_ref_.data(), d_ref, len, cudaMemcpyDeviceToHost);
  }

  explicit device_equals_t(std::vector<DataT> ref) : h_ref_(std::move(ref)) {}

  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_ref_.size(); ++i) {
      if (std::isnan(h_ref_[i]) && std::isnan(h_computed[i])) continue;
      if (h_ref_[i] != h_computed[i]) {
        UNSCOPED_INFO("@i=" << i << ": " << h_computed[i] << " != " << h_ref_[i]
                            << " (val != ref)");
        return false;
      }
    }
    return true;
  }

  [[nodiscard]] std::string describe() const override { return "device array matcher"; }
};  // class device_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
 *
 * @return the matcher object
 */
template <typename DataT>
device_equals_t<DataT> device_equals(DataT* d_ref, size_t len)
{
  return device_equals_t<DataT>(d_ref, len);
}

/**
 * @brief Builder function for the custom device array matcher
 *
 * @tparam DataT data type of the arrays
 *
 * @param[in] ref reference host vector
 *
 * @return the matcher object
 */
template <typename DataT>
device_equals_t<DataT> device_equals(const std::vector<DataT>& ref)
{
  return device_equals_t<DataT>(ref);
}

}  // namespace cugraph::ops::catch2
