/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

// Copyright 2004-present Facebook. All Rights Reserved.
// -*- c++ -*-

#include "IndexIVF_c.h"
#include <faiss/IndexIVF.h>
#include "Clustering_c.h"
#include "Index_c.h"
#include "macros_impl.h"

using faiss::IndexIVF;
using faiss::IndexIVFStats;

DEFINE_DESTRUCTOR(IndexIVF)
DEFINE_INDEX_DOWNCAST(IndexIVF)

/// number of possible key values
DEFINE_GETTER(IndexIVF, size_t, nlist)
/// number of probes at query time
DEFINE_GETTER(IndexIVF, size_t, nprobe)
DEFINE_SETTER(IndexIVF, size_t, nprobe)

/// quantizer that maps vectors to inverted lists
DEFINE_GETTER_PERMISSIVE(IndexIVF, FaissIndex*, quantizer)

/**
 * = 0: use the quantizer as index in a kmeans training
 * = 1: just pass on the training set to the train() of the quantizer
 * = 2: kmeans training on a flat index + add the centroids to the quantizer
 */
DEFINE_GETTER(IndexIVF, char, quantizer_trains_alone)

/// whether object owns the quantizer
DEFINE_GETTER(IndexIVF, int, own_fields)
DEFINE_SETTER(IndexIVF, int, own_fields)

using faiss::IndexIVF;

int faiss_IndexIVF_merge_from(
        FaissIndexIVF* index,
        FaissIndexIVF* other,
        idx_t add_id) {
    try {
        reinterpret_cast<IndexIVF*>(index)->merge_from(
                *reinterpret_cast<IndexIVF*>(other), add_id);
    }
    CATCH_AND_HANDLE
}

int faiss_IndexIVF_copy_subset_to(
        const FaissIndexIVF* index,
        FaissIndexIVF* other,
        int subset_type,
        idx_t a1,
        idx_t a2) {
    try {
        reinterpret_cast<const IndexIVF*>(index)->copy_subset_to(
                *reinterpret_cast<IndexIVF*>(other), subset_type, a1, a2);
    }
    CATCH_AND_HANDLE
}

int faiss_IndexIVF_search_preassigned(
        const FaissIndexIVF* index,
        idx_t n,
        const float* x,
        idx_t k,
        const idx_t* assign,
        const float* centroid_dis,
        float* distances,
        idx_t* labels,
        int store_pairs) {
    try {
        reinterpret_cast<const IndexIVF*>(index)->search_preassigned(
                n, x, k, assign, centroid_dis, distances, labels, store_pairs);
    }
    CATCH_AND_HANDLE
}

size_t faiss_IndexIVF_get_list_size(
        const FaissIndexIVF* index,
        size_t list_no) {
    return reinterpret_cast<const IndexIVF*>(index)->get_list_size(list_no);
}

int faiss_IndexIVF_make_direct_map(
        FaissIndexIVF* index,
        int new_maintain_direct_map) {
    try {
        reinterpret_cast<IndexIVF*>(index)->make_direct_map(
                static_cast<bool>(new_maintain_direct_map));
    }
    CATCH_AND_HANDLE
}

double faiss_IndexIVF_imbalance_factor(const FaissIndexIVF* index) {
    return reinterpret_cast<const IndexIVF*>(index)
            ->invlists->imbalance_factor();
}

/// display some stats about the inverted lists
void faiss_IndexIVF_print_stats(const FaissIndexIVF* index) {
    reinterpret_cast<const IndexIVF*>(index)->invlists->print_stats();
}

/// get inverted lists ids
void faiss_IndexIVF_invlists_get_ids(
        const FaissIndexIVF* index,
        size_t list_no,
        idx_t* invlist) {
    const idx_t* list =
            reinterpret_cast<const IndexIVF*>(index)->invlists->get_ids(
                    list_no);
    size_t list_size =
            reinterpret_cast<const IndexIVF*>(index)->get_list_size(list_no);
    memcpy(invlist, list, list_size * sizeof(idx_t));
}

void faiss_IndexIVFStats_reset(FaissIndexIVFStats* stats) {
    reinterpret_cast<IndexIVFStats*>(stats)->reset();
}

FaissIndexIVFStats* faiss_get_indexIVF_stats() {
    return reinterpret_cast<FaissIndexIVFStats*>(&faiss::indexIVF_stats);
}