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

#include <catch2/catch.hpp>

#include <vector>

namespace cugraph::ops::catch2 {

/**
 * @brief Custom matcher class for matching edge types
 *        according to sampled edges
 *
 * @tparam DataT data type of the arrays
 */
template <typename DataT>
class device_sampled_etypes_t : public Catch::MatcherBase<DataT*> {
  std::vector<DataT> h_neighbors_;
  std::vector<DataT> h_dst_nodes_;
  std::vector<DataT> h_offsets_;
  std::vector<DataT> h_indices_;
  std::vector<DataT> h_full_etypes_;
  const DataT sample_size_;

 public:
  explicit device_sampled_etypes_t(DataT* d_neighbors,
                                   DataT* d_dst_nodes,
                                   DataT n_dst_nodes,
                                   DataT sample_size,
                                   DataT* offsets,
                                   size_t n_offsets,
                                   DataT* indices,
                                   size_t n_indices,
                                   DataT* full_edge_types)
    : sample_size_(sample_size)
  {
    auto neigh_len = static_cast<size_t>(n_dst_nodes) * static_cast<size_t>(sample_size);
    h_neighbors_.resize(neigh_len);
    utils::copy(h_neighbors_.data(), d_neighbors, neigh_len, cudaMemcpyDeviceToHost);
    if (d_dst_nodes != nullptr) {
      h_dst_nodes_.resize(n_dst_nodes);
      utils::copy(h_dst_nodes_.data(), d_dst_nodes, n_dst_nodes, cudaMemcpyDeviceToHost);
    }
    h_offsets_.resize(n_offsets);
    utils::copy(h_offsets_.data(), offsets, n_offsets, cudaMemcpyDeviceToHost);
    h_indices_.resize(n_indices);
    utils::copy(h_indices_.data(), indices, n_indices, cudaMemcpyDeviceToHost);
    h_full_etypes_.resize(n_indices);
    utils::copy(h_full_etypes_.data(), full_edge_types, n_indices, cudaMemcpyDeviceToHost);
  }

  bool match(DataT* const& d_etypes) const override
  {
    std::vector<DataT> h_etypes(h_neighbors_.size());
    utils::copy(h_etypes.data(), d_etypes, h_neighbors_.size(), cudaMemcpyDeviceToHost);
    for (size_t i = 0; i < h_neighbors_.size(); ++i) {
      auto etype = h_etypes[i];
      auto n_id  = h_neighbors_[i];
      if (n_id == graph::INVALID_ID<DataT> && etype != graph::DEFAULT_EDGE_TYPE) {
        UNSCOPED_INFO("@i=" << i << " invalid neighbor, but etype=" << etype);
        return false;
      }
      if (n_id == graph::INVALID_ID<DataT>) continue;
      auto out_n  = i / sample_size_;
      auto out_id = h_dst_nodes_.empty() ? out_n : h_dst_nodes_[out_n];
      // TODO(mjoux) improve this once we pass edge indices along in sampling
      auto found_edge = false;
      for (auto j = h_offsets_[out_id]; j < h_offsets_[out_id + 1]; ++j) {
        if (h_indices_[j] == n_id && h_full_etypes_[j] == etype) found_edge = true;
      }
      if (!found_edge) {
        UNSCOPED_INFO("@i=" << i << ", out node=" << out_n << ", out id=" << out_id << ", neighbor="
                            << n_id << " : no edge with type " << etype << " found.");
        return false;
      }
    }
    return true;
  }

  [[nodiscard]] std::string describe() const override
  {
    return "device sampled edge types matcher";
  }
};  // class device_sampled_etypes_t

/**
 * @brief Custom matcher class for matching edge types
 *        according to a range of expected types
 *
 * @tparam DataT data type of the arrays
 */
template <typename DataT>
class device_sampled_etypes_range_t : public Catch::MatcherBase<DataT*> {
  std::vector<DataT> h_neighbors_;
  DataT start_, end_;

 public:
  explicit device_sampled_etypes_range_t(DataT* d_neighbors,
                                         DataT neigh_len,
                                         DataT _start,
                                         DataT _end)
    : start_(_start), end_(_end)
  {
    h_neighbors_.resize(static_cast<size_t>(neigh_len));
    utils::copy(
      h_neighbors_.data(), d_neighbors, static_cast<size_t>(neigh_len), cudaMemcpyDeviceToHost);
  }

  bool match(DataT* const& d_etypes) const override
  {
    std::vector<DataT> h_etypes(h_neighbors_.size());
    utils::copy(h_etypes.data(), d_etypes, h_neighbors_.size(), cudaMemcpyDeviceToHost);
    for (size_t i = 0; i < h_neighbors_.size(); ++i) {
      auto etype = h_etypes[i];
      auto n_id  = h_neighbors_[i];
      if (n_id == graph::INVALID_ID<DataT>) continue;
      if (etype < start_ || etype >= end_) {
        UNSCOPED_INFO("@i=" << i << ", etype=" << etype << " not in range [" << start_ << ", "
                            << end_ << "].");
        return false;
      }
    }
    return true;
  }

  [[nodiscard]] std::string describe() const override
  {
    return "device sampled edge types range matcher";
  }
};  // class device_sampled_etypes_range_t

/**
 * @brief Builder function for the custom device sampled edge types matcher
 *
 * @tparam DataT data type of the arrays
 *
 * @param[in]   d_neighbors     sampled neighbors [device]
 *                              [n_dst_nodes x sample_size]
 * @param[in]   d_dst_nodes     sampled out nodes (map to full graph nodes)
 *                              [device] [n_dst_nodes]
 * @param[in]   n_dst_nodes     number of output nodes that were sampled
 * @param[in]   sample_size     max number of neighbors per output node
 * @param[in]   offsets         full graph CSR offsets [device] [n_offsets]
 * @param[in]   n_offsets       number of offsets
 * @param[in]   indices         full graph CSR indices [device] [n_indices]
 * @param[in]   n_indices       number of indices
 * @param[in]   full_edge_types full graph edge types [device] [n_indices]
 *
 * @return the matcher object
 */
template <typename DataT>
device_sampled_etypes_t<DataT> device_sampled_edge_types(DataT* d_neighbors,
                                                         DataT* d_dst_nodes,
                                                         DataT n_dst_nodes,
                                                         DataT sample_size,
                                                         DataT* offsets,
                                                         size_t n_offsets,
                                                         DataT* indices,
                                                         size_t n_indices,
                                                         DataT* full_edge_types)
{
  return device_sampled_etypes_t<DataT>(d_neighbors,
                                        d_dst_nodes,
                                        n_dst_nodes,
                                        sample_size,
                                        offsets,
                                        n_offsets,
                                        indices,
                                        n_indices,
                                        full_edge_types);
}

/**
 * @brief Builder for the custom device sampled edge types range matcher
 *
 * @tparam DataT data type of the arrays
 *
 * @param[in]   d_neighbors     sampled neighbors [device] [neigh_len]
 * @param[in]   neigh_len       number of sampled neighbors
 * @param[in]   start           start of valid edge type range
 * @param[in]   end             end of valid edge type range (not inclusive)
 *
 * @return the matcher object
 */
template <typename DataT>
device_sampled_etypes_range_t<DataT> device_sampled_edge_types_range(DataT* d_neighbors,
                                                                     DataT neigh_len,
                                                                     DataT start,
                                                                     DataT end)
{
  return device_sampled_etypes_range_t<DataT>(d_neighbors, neigh_len, start, end);
}

}  // namespace cugraph::ops::catch2
