

#pragma once

#include "cuda_stub.hpp"  // cudaStream_t

#include <stdexcept>
#include <string>
#include <vector>

namespace cuvs::bench {

enum Objective {
  THROUGHPUT,  // See how many vectors we can push through
  LATENCY      // See how fast we can push a vector through
};

enum class MemoryType {
  Host,
  HostMmap,
  Device,
};

enum class Metric {
  kInnerProduct,
  kEuclidean,
};

inline auto parse_metric(const std::string& metric_str) -> Metric
{
  if (metric_str == "inner_product") {
    return cuvs::bench::Metric::kInnerProduct;
  } else if (metric_str == "euclidean") {
    return cuvs::bench::Metric::kEuclidean;
  } else {
    throw std::runtime_error("invalid metric: '" + metric_str + "'");
  }
}

inline auto parse_memory_type(const std::string& memory_type) -> MemoryType
{
  if (memory_type == "host") {
    return MemoryType::Host;
  } else if (memory_type == "mmap") {
    return MemoryType::HostMmap;
  } else if (memory_type == "device") {
    return MemoryType::Device;
  } else {
    throw std::runtime_error("invalid memory type: '" + memory_type + "'");
  }
}

class AlgoProperty {
 public:
  inline AlgoProperty() {}
  inline AlgoProperty(MemoryType dataset_memory_type_, MemoryType query_memory_type_)
    : dataset_memory_type(dataset_memory_type_), query_memory_type(query_memory_type_)
  {
  }
  MemoryType dataset_memory_type;
  // neighbors/distances should have same memory type as queries
  MemoryType query_memory_type;
  virtual ~AlgoProperty() = default;
};

class AnnBase {
 public:
  inline AnnBase(Metric metric, int dim) : metric_(metric), dim_(dim) {}
  virtual ~AnnBase() = default;

 protected:
  Metric metric_;
  int dim_;
};

template <typename T>
class ANN : public AnnBase {
 public:
  struct AnnSearchParam {
    Objective metric_objective = Objective::LATENCY;
    virtual ~AnnSearchParam()  = default;
    [[nodiscard]] virtual auto needs_dataset() const -> bool { return false; };
  };

  inline ANN(Metric metric, int dim) : AnnBase(metric, dim) {}

  virtual void build(const T* dataset, size_t nrow, cudaStream_t stream = 0) = 0;

  virtual void set_search_param(const AnnSearchParam& param) = 0;
  // TODO: this assumes that an algorithm can always return k results.
  // This is not always possible.
  virtual void search(const T* queries,
                      int batch_size,
                      int k,
                      size_t* neighbors,
                      float* distances,
                      cudaStream_t stream = 0) const = 0;

  virtual void save(const std::string& file) const = 0;
  virtual void load(const std::string& file)       = 0;

  virtual AlgoProperty get_preference() const = 0;

  // Some algorithms don't save the building dataset in their indices.
  // So they should be given the access to that dataset during searching.
  // The advantage of this way is that index has smaller size
  // and many indices can share one dataset.
  //
  // SearchParam::needs_dataset() of such algorithm should be true,
  // and set_search_dataset() should save the passed-in pointer somewhere.
  // The client code should call set_search_dataset() before searching,
  // and should not release dataset before searching is finished.
  virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/){};
};

}  // namespace cuvs::bench

#define REGISTER_ALGO_INSTANCE(DataT)                                                              \
  template auto cuvs::bench::create_algo<DataT>(                                                   \
    const std::string&, const std::string&, int, const nlohmann::json&, const std::vector<int>&)   \
    ->std::unique_ptr<cuvs::bench::ANN<DataT>>;                                                    \
  template auto cuvs::bench::create_search_param<DataT>(const std::string&, const nlohmann::json&) \
    ->std::unique_ptr<typename cuvs::bench::ANN<DataT>::AnnSearchParam>;
