

#pragma once

#include <cstdint>  // int64_t

#include <cuvs/distance/distance_types.hpp>  // cuvs::distance::DistanceType
#include <raft/core/host_mdspan.hpp>         // raft::host_matrix_view
#include <raft/util/raft_explicit.hpp>       // RAFT_EXPLICIT

#ifdef RAFT_EXPLICIT_INSTANTIATE_ONLY

namespace cuvs::neighbors::detail {

template <typename IdxT, typename DataT, typename DistanceT, typename ExtentsT>
[[gnu::optimize(3), gnu::optimize("tree-vectorize")]] void refine_host(
  raft::host_matrix_view<const DataT, ExtentsT, raft::row_major> dataset,
  raft::host_matrix_view<const DataT, ExtentsT, raft::row_major> queries,
  raft::host_matrix_view<const IdxT, ExtentsT, raft::row_major> neighbor_candidates,
  raft::host_matrix_view<IdxT, ExtentsT, raft::row_major> indices,
  raft::host_matrix_view<DistanceT, ExtentsT, raft::row_major> distances,
  distance::DistanceType metric = distance::DistanceType::L2Unexpanded) RAFT_EXPLICIT;

}

#endif  // RAFT_EXPLICIT_INSTANTIATE_ONLY

#define instantiate_raft_neighbors_refine(IdxT, DataT, DistanceT, ExtentsT)                    \
  extern template void cuvs::neighbors::detail::refine_host<IdxT, DataT, DistanceT, ExtentsT>( \
    raft::host_matrix_view<const DataT, ExtentsT, raft::row_major> dataset,                    \
    raft::host_matrix_view<const DataT, ExtentsT, raft::row_major> queries,                    \
    raft::host_matrix_view<const IdxT, ExtentsT, raft::row_major> neighbor_candidates,         \
    raft::host_matrix_view<IdxT, ExtentsT, raft::row_major> indices,                           \
    raft::host_matrix_view<DistanceT, ExtentsT, raft::row_major> distances,                    \
    distance::DistanceType metric);

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

#undef instantiate_raft_neighbors_refine
