

#pragma once

#include <cuvs/neighbors/ivf_flat_types.hpp>
#include <string>

namespace raft::runtime::neighbors::ivf_flat {

// We define overloads for build and extend with void return type. This is used in the Cython
// wrappers, where exception handling is not compatible with return type that has nontrivial
// constructor.
#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                              \
  auto build(raft::resources const& handle,                                          \
             const cuvs::neighbors::ivf_flat::index_params& params,                  \
             raft::device_matrix_view<const T, IdxT, row_major> dataset)             \
    ->cuvs::neighbors::ivf_flat::index<T, IdxT>;                                     \
                                                                                     \
  auto extend(raft::resources const& handle,                                         \
              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,        \
              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, \
              const cuvs::neighbors::ivf_flat::index<T, IdxT>& orig_index)           \
    ->cuvs::neighbors::ivf_flat::index<T, IdxT>;                                     \
                                                                                     \
  void build(raft::resources const& handle,                                          \
             const cuvs::neighbors::ivf_flat::index_params& params,                  \
             raft::device_matrix_view<const T, IdxT, row_major> dataset,             \
             cuvs::neighbors::ivf_flat::index<T, IdxT>& idx);                        \
                                                                                     \
  void extend(raft::resources const& handle,                                         \
              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,        \
              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, \
              cuvs::neighbors::ivf_flat::index<T, IdxT>* idx);                       \
                                                                                     \
  void serialize_file(raft::resources const& handle,                                 \
                      const std::string& filename,                                   \
                      const cuvs::neighbors::ivf_flat::index<T, IdxT>& index);       \
                                                                                     \
  void deserialize_file(raft::resources const& handle,                               \
                        const std::string& filename,                                 \
                        cuvs::neighbors::ivf_flat::index<T, IdxT>* index);           \
  void serialize(raft::resources const& handle,                                      \
                 std::string& str,                                                   \
                 const cuvs::neighbors::ivf_flat::index<T, IdxT>& index);            \
  void deserialize(raft::resources const& handle,                                    \
                   const std::string& str,                                           \
                   cuvs::neighbors::ivf_flat::index<T, IdxT>*);

RAFT_INST_BUILD_EXTEND(float, int64_t)
RAFT_INST_BUILD_EXTEND(int8_t, int64_t)
RAFT_INST_BUILD_EXTEND(uint8_t, int64_t)

#undef RAFT_INST_BUILD_EXTEND

#define RAFT_INST_SEARCH(T, IdxT)                                 \
  void search(raft::resources const&,                             \
              cuvs::neighbors::ivf_flat::search_params const&,    \
              cuvs::neighbors::ivf_flat::index<T, IdxT> const&,   \
              raft::device_matrix_view<const T, IdxT, row_major>, \
              raft::device_matrix_view<IdxT, IdxT, row_major>,    \
              raft::device_matrix_view<float, IdxT, row_major>);

RAFT_INST_SEARCH(float, int64_t);
RAFT_INST_SEARCH(int8_t, int64_t);
RAFT_INST_SEARCH(uint8_t, int64_t);

#undef RAFT_INST_SEARCH

}  // namespace raft::runtime::neighbors::ivf_flat
