
#pragma once

#include <cassert>
#include <cuvs/distance/detail/distance.cuh>
#include <cuvs/distance/distance_types.hpp>
#include <cuvs/spatial/knn/detail/fused_l2_knn.cuh>
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>

#include "../common/ann_types.hpp"

namespace raft_temp {

inline cuvs::distance::DistanceType parse_metric_type(cuvs::bench::Metric metric)
{
  if (metric == cuvs::bench::Metric::kInnerProduct) {
    return cuvs::distance::DistanceType::InnerProduct;
  } else if (metric == cuvs::bench::Metric::kEuclidean) {
    return cuvs::distance::DistanceType::L2Expanded;
  } else {
    throw std::runtime_error("raft supports only metric type of inner product and L2");
  }
}

}  // namespace raft_temp

namespace cuvs::bench {

// brute force fused L2 KNN - RAFT
template <typename T>
class RaftGpu : public ANN<T> {
 public:
  using typename ANN<T>::AnnSearchParam;

  RaftGpu(Metric metric, int dim);

  void build(const T*, size_t, cudaStream_t) final;

  void set_search_param(const AnnSearchParam& param) override;

  // TODO: if the number of results is less than k, the remaining elements of 'neighbors'
  // will be filled with (size_t)-1
  void search(const T* queries,
              int batch_size,
              int k,
              size_t* neighbors,
              float* distances,
              cudaStream_t stream = 0) const final;

  // to enable dataset access from GPU memory
  AlgoProperty get_preference() const override
  {
    AlgoProperty property;
    property.dataset_memory_type = MemoryType::Device;
    property.query_memory_type   = MemoryType::Device;
    return property;
  }
  void set_search_dataset(const T* dataset, size_t nrow) override;
  void save(const std::string& file) const override;
  void load(const std::string&) override { return; };

 protected:
  cuvs::distance::DistanceType metric_type_;
  int device_;
  const T* dataset_;
  size_t nrow_;
};

template <typename T>
RaftGpu<T>::RaftGpu(Metric metric, int dim)
  : ANN<T>(metric, dim), metric_type_(raft_temp::parse_metric_type(metric))
{
  static_assert(std::is_same_v<T, float>, "raft support only float type");
  assert(metric_type_ == cuvs::distance::DistanceType::L2Expanded);
  RAFT_CUDA_TRY(cudaGetDevice(&device_));
}

template <typename T>
void RaftGpu<T>::build(const T*, size_t, cudaStream_t)
{
  // as this is brute force algo so no index building required
  return;
}

template <typename T>
void RaftGpu<T>::set_search_param(const AnnSearchParam&)
{
  // Nothing to set here as it is brute force implementation
}

template <typename T>
void RaftGpu<T>::set_search_dataset(const T* dataset, size_t nrow)
{
  dataset_ = dataset;
  nrow_    = nrow;
}

template <typename T>
void RaftGpu<T>::save(const std::string& file) const
{
  // create a empty index file as no index to store.
  std::fstream fp;
  fp.open(file.c_str(), std::ios::out);
  if (!fp) {
    printf("Error in creating file!!!\n");
    ;
    return;
  }
  fp.close();
}

template <typename T>
void RaftGpu<T>::search(const T* queries,
                        int batch_size,
                        int k,
                        size_t* neighbors,
                        float* distances,
                        cudaStream_t stream) const
{
  // TODO: Integrate new `raft::brute_force::index` (from
  // https://github.com/rapidsai/raft/pull/1817)
  cuvs::spatial::knn::detail::fusedL2Knn(this->dim_,
                                         reinterpret_cast<int64_t*>(neighbors),
                                         distances,
                                         dataset_,
                                         queries,
                                         nrow_,
                                         static_cast<size_t>(batch_size),
                                         k,
                                         true,
                                         true,
                                         stream,
                                         metric_type_);
}

}  // namespace cuvs::bench
