/*
 * Copyright (c) 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.
 */

#include "device_approx_equals.hpp"

#include <catch2/catch.hpp>

#include <sstream>
#include <utility>

namespace cugraph::ops::catch2 {

// defined in tests/src/catch2/mean_var.cu[h]
template <typename DataT>
std::pair<DataT, DataT> eval_mean_var(const DataT* data, int len);

/**
 * @brief Custom matcher class for comparing mean and variance of an array
 *
 * @tparam DataT data type of the arrays
 */
template <typename DataT>
class mean_var_t : public Catch::MatcherBase<std::pair<DataT, DataT>> {
  compare_approx<DataT> matcher_;
  std::pair<DataT, DataT> ref_;
  const DataT tolerance_;

 public:
  mean_var_t(DataT mean, DataT var, DataT tol) : matcher_(tol), tolerance_(tol)
  {
    ref_.first  = mean;
    ref_.second = var;
  }

  bool match(std::pair<DataT, DataT> const& computed) const override
  {
    return matcher_(ref_.first, computed.first) && matcher_(ref_.second, computed.second);
  }

  [[nodiscard]] std::string describe() const override
  {
    std::stringstream ss;
    ss << "ref mean,var=" << ref_.first << "," << ref_.second << " tolerance=" << tolerance_;
    return ss.str();
  }
};  // class mean_var_t

/**
 * @brief Builder function for the mean/variance matcher
 *
 * @tparam DataT data type of the arrays
 *
 * @param[in] mean reference mean value
 * @param[in] var  reference variance value
 * @param[in] tol  tolerance for comparison
 *
 * @return the matcher object
 */
template <typename DataT>
mean_var_t<DataT> mean_var(DataT mean, DataT var, DataT tol)
{
  return mean_var_t<DataT>(mean, var, tol);
}

}  // namespace cugraph::ops::catch2
