/*
 * 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 checking an undirected CSR
 *
 * @tparam IdxT index type of the graph
 */
template <typename IdxT>
class device_csc_undirected_t : public Catch::MatcherBase<IdxT*> {
  std::vector<IdxT> h_offsets_;
  std::vector<IdxT> h_indices_;
  size_t n_dst_nodes_;

 public:
  explicit device_csc_undirected_t(IdxT* csc_offsets, IdxT* csc_indices, IdxT n_dst_nodes, IdxT nnz)
  {
    h_offsets_.resize(n_dst_nodes + 1);
    n_dst_nodes_ = h_offsets_.size() - 1;
    utils::copy(h_offsets_.data(), csc_offsets, h_offsets_.size(), cudaMemcpyDeviceToHost);
    h_indices_.resize(nnz);
    utils::copy(h_indices_.data(), csc_indices, h_indices_.size(), cudaMemcpyDeviceToHost);
  }

  bool match(IdxT* const& csc_rev_idx) const override
  {
    std::vector<IdxT> h_rev_idx(h_indices_.size());
    utils::copy(h_rev_idx.data(), csc_rev_idx, h_rev_idx.size(), cudaMemcpyDeviceToHost);

    for (size_t dst_id = 0; dst_id < n_dst_nodes_; ++dst_id) {
      auto off_start = h_offsets_[dst_id], off_end = h_offsets_[dst_id + 1];
      for (auto i = off_start; i < off_end; ++i) {
        auto rev_edge = h_rev_idx[i];
        if (rev_edge < 0 || rev_edge >= h_indices_.size()) {
          UNSCOPED_INFO("@node=" << dst_id << " / edge=" << i << " invalid reverse edge ID "
                                 << rev_edge << " (<0 or >" << h_indices_.size() << ")");
          return false;
        }
        auto rev_src_id = h_indices_[rev_edge];
        if (dst_id != rev_src_id) {
          UNSCOPED_INFO("@node=" << dst_id << " / edge=" << i << " reverse edge " << rev_edge
                                 << " has node=" << rev_src_id);
          return false;
        }
        // binary search in offsets to get reverse edge destination node
        auto b1 = size_t{0}, b2 = n_dst_nodes_;
        auto rev_dst_id = graph::INVALID_ID<IdxT>;
        while (true) {
          auto mid = (b1 + b2) / 2;
          if (mid == b1 && mid == b2) break;
          auto o1 = h_offsets_[mid], o2 = h_offsets_[mid + 1];
          if (rev_edge < o1) {
            b2 = mid;
          } else if (rev_edge >= o2) {
            b1 = mid;
          } else {
            rev_dst_id = mid;
            break;
          }
        }
        if (rev_dst_id == graph::INVALID_ID<IdxT>) {
          UNSCOPED_INFO("@node=" << dst_id << " / edge=" << i << " reverse edge " << rev_edge
                                 << " unable to find dst node");
          return false;
        }
        auto src_id = h_indices_[i];
        if (rev_dst_id != src_id) {
          UNSCOPED_INFO("@node=" << dst_id << " / edge=" << i << " reverse edge " << rev_edge
                                 << " has dst ID " << rev_dst_id << " but expected src ID "
                                 << src_id);
          return false;
        }
      }
    }
    return true;
  }

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

/**
 * @brief Builder function for the custom undirected CSR matcher
 *
 * @note It is expected that `csc_offsets` and `csc_indices` have been checked
 *       before using this matcher: they should have the expected sizes and
 *       offsets should be monotonic (with values in [0, nnz]) and indices
 *       should have values in [0, n_dst_nodes - 1]
 *
 * @tparam IdxT index type of the MFGs
 *
 * @param[in] csc_offsets   reference CSR offsets
 * @param[in] csc_indices   reference CSR indices
 * @param[in] n_dst_nodes       number of nodes (1 less than size of `csc_offsets`)
 * @param[in] nnz           number of edges (size of `csc_indices`)
 *
 * @return the matcher object
 */
template <typename IdxT>
device_csc_undirected_t<IdxT> device_csc_undirected(IdxT* csc_offsets,
                                                    IdxT* csc_indices,
                                                    IdxT n_dst_nodes,
                                                    IdxT nnz)
{
  return device_csc_undirected_t<IdxT>(csc_offsets, csc_indices, n_dst_nodes, nnz);
}

}  // namespace cugraph::ops::catch2
