/*
 * 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 <cugraph-ops/graph/format.hpp>

#include <catch2/catch.hpp>

#include <vector>

namespace cugraph::ops::catch2 {

/**
 * @brief Custom matcher class for comparing CSR vs ELLPACK graphs
 *
 * @tparam IdxT index type of the MFGs
 */
template <typename IdxT>
class device_csr_ellpack_equals_t : public Catch::MatcherBase<graph::mfg_csr<IdxT>> {
  graph::mfg_ellpack<IdxT> mfg_ellpack_;
  std::vector<IdxT> ellpack_neighbors_;
  std::vector<IdxT> ellpack_counts_;
  IdxT n_total_neighbors_;

 public:
  explicit device_csr_ellpack_equals_t(graph::mfg_ellpack<IdxT>& mfg)
    : mfg_ellpack_(mfg), n_total_neighbors_(0)
  {
    ellpack_neighbors_.resize(mfg.n_out_nodes * mfg.sample_size);
    utils::copy(
      ellpack_neighbors_.data(), mfg.neighbors, ellpack_neighbors_.size(), cudaMemcpyDeviceToHost);
    ellpack_counts_.resize(mfg.n_out_nodes);
    if (mfg.neighbor_counts == nullptr) {
      for (IdxT i = 0; i < mfg.n_out_nodes; ++i) {
        for (IdxT s = 0; s < mfg.sample_size; ++s) {
          if (ellpack_neighbors_[i * mfg.sample_size + s] == graph::INVALID_ID<IdxT>) continue;
          ellpack_counts_[i]++;
        }
      }
    } else {
      utils::copy(ellpack_counts_.data(),
                  mfg.neighbor_counts,
                  ellpack_counts_.size(),
                  cudaMemcpyDeviceToHost);
    }
    for (auto c : ellpack_counts_)
      n_total_neighbors_ += c;
  }

  bool match(graph::mfg_csr<IdxT> const& mfg_computed) const override
  {
    // first check if the number of offsets matches
    if (mfg_ellpack_.n_out_nodes != mfg_computed.n_out_nodes) {
      UNSCOPED_INFO("Got " << mfg_computed.n_out_nodes << " out nodes, but expected "
                           << mfg_ellpack_.n_out_nodes);
      return false;
    }
    // check the number of indices
    if (n_total_neighbors_ != mfg_computed.n_indices) {
      UNSCOPED_INFO("Got " << mfg_computed.n_indices << " neighbors, but expected "
                           << n_total_neighbors_);
      return false;
    }
    // check that all counts match
    std::vector<IdxT> csr_offsets(ellpack_counts_.size() + 1);
    utils::copy(
      csr_offsets.data(), mfg_computed.offsets, csr_offsets.size(), cudaMemcpyDeviceToHost);
    for (size_t i = 0; i < csr_offsets.size() - 1; ++i) {
      auto n_expected = ellpack_counts_[i];
      auto n_computed = csr_offsets[i + 1] - csr_offsets[i];
      if (n_expected != n_computed) {
        UNSCOPED_INFO("@i=" << i << " expected " << n_expected << " neighbors, but got "
                            << n_computed);
        return false;
      }
    }

    // check that all indices match
    std::vector<IdxT> csr_indices(n_total_neighbors_);
    utils::copy(
      csr_indices.data(), mfg_computed.indices, csr_indices.size(), cudaMemcpyDeviceToHost);

    for (size_t i = 0; i < csr_offsets.size() - 1; ++i) {
      auto off_start = csr_offsets[i], off_end = csr_offsets[i + 1];
      for (auto n = off_start; n < off_end; ++n) {
        if (n < 0 || n >= static_cast<IdxT>(csr_indices.size())) {
          UNSCOPED_INFO("@i=" << i << ": offset " << n << " out of range of indices array "
                              << csr_indices.size());
          return false;
        }
        auto nidx          = csr_indices[n];
        auto i_ellpack     = i * static_cast<size_t>(mfg_ellpack_.sample_size) + n - off_start;
        auto nidx_expected = ellpack_neighbors_[i_ellpack];
        if (nidx_expected != nidx) {
          UNSCOPED_INFO("@i=" << i << ", n=" << n - off_start << " expected neighbor idx "
                              << nidx_expected << " but got " << nidx);
          return false;
        }
      }
    }

    return true;
  }

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

/**
 * @brief Builder function for the custom mfg matcher
 *
 * @tparam IdxT index type of the MFGs
 *
 * @param[in] mfg   reference ELLPACK MFG
 *
 * @return the matcher object
 */
template <typename IdxT>
device_csr_ellpack_equals_t<IdxT> device_csr_ellpack_equals(graph::mfg_ellpack<IdxT>& mfg)
{
  return device_csr_ellpack_equals_t<IdxT>(mfg);
}

}  // namespace cugraph::ops::catch2
