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

// -*- c++ -*-

#include <faiss/IndexIVFPQR.h>

#include <cinttypes>

#include <faiss/utils/Heap.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/utils.h>

#include <faiss/impl/FaissAssert.h>

namespace faiss {

/*****************************************
 * IndexIVFPQR implementation
 ******************************************/

IndexIVFPQR::IndexIVFPQR(
        Index* quantizer,
        size_t d,
        size_t nlist,
        size_t M,
        size_t nbits_per_idx,
        size_t M_refine,
        size_t nbits_per_idx_refine)
        : IndexIVFPQ(quantizer, d, nlist, M, nbits_per_idx),
          refine_pq(d, M_refine, nbits_per_idx_refine),
          k_factor(4) {
    by_residual = true;
    refine_pq.cp.max_points_per_centroid = 1000;
}

IndexIVFPQR::IndexIVFPQR() : k_factor(1) {
    by_residual = true;
    refine_pq.cp.max_points_per_centroid = 1000;
}

void IndexIVFPQR::reset() {
    IndexIVFPQ::reset();
    refine_codes.clear();
}

void IndexIVFPQR::train_encoder(idx_t n, const float* x, const idx_t* assign) {
    IndexIVFPQ::train_encoder(n, x, assign);
    if (verbose) {
        printf("training %zdx%zd 2nd level PQ quantizer on %" PRId64
               " %dD-vectors\n",
               refine_pq.M,
               refine_pq.ksub,
               n,
               d);
    }
    refine_pq.cp.verbose = verbose;

    // 2nd level residual
    std::vector<float> residual_2(n * d);
    std::vector<uint8_t> train_codes(pq.code_size * n);
    pq.compute_codes(x, train_codes.data(), n);

    for (idx_t i = 0; i < n; i++) {
        const float* xx = x + i * d;
        float* res = residual_2.data() + i * d;
        pq.decode(train_codes.data() + i * pq.code_size, res);
        for (int j = 0; j < d; j++) {
            res[j] = xx[j] - res[j];
        }
    }

    refine_pq.train(n, residual_2.data());
}

idx_t IndexIVFPQR::train_encoder_num_vectors() const {
    return std::max(
            pq.cp.max_points_per_centroid * pq.ksub,
            refine_pq.cp.max_points_per_centroid * refine_pq.ksub);
}

void IndexIVFPQR::add_with_ids(idx_t n, const float* x, const idx_t* xids) {
    add_core(n, x, xids, nullptr);
}

void IndexIVFPQR::add_core(
        idx_t n,
        const float* x,
        const idx_t* xids,
        const idx_t* precomputed_idx,
        void* /*inverted_list_context*/) {
    std::unique_ptr<float[]> residual_2(new float[n * d]);

    idx_t n0 = ntotal;

    add_core_o(n, x, xids, residual_2.get(), precomputed_idx);

    refine_codes.resize(ntotal * refine_pq.code_size);

    refine_pq.compute_codes(
            residual_2.get(), &refine_codes[n0 * refine_pq.code_size], n);
}
#define TIC t0 = get_cycles()
#define TOC get_cycles() - t0

void IndexIVFPQR::search_preassigned(
        idx_t n,
        const float* x,
        idx_t k,
        const idx_t* idx,
        const float* L1_dis,
        float* distances,
        idx_t* labels,
        bool store_pairs,
        const IVFSearchParameters* params,
        IndexIVFStats* stats) const {
    uint64_t t0;
    TIC;
    size_t k_coarse = long(k * k_factor);
    std::unique_ptr<idx_t[]> coarse_labels(new idx_t[k_coarse * n]);
    {
        // query with quantizer levels 1 and 2.
        std::unique_ptr<float[]> coarse_distances(new float[k_coarse * n]);

        IndexIVFPQ::search_preassigned(
                n,
                x,
                k_coarse,
                idx,
                L1_dis,
                coarse_distances.get(),
                coarse_labels.get(),
                true,
                params);
    }

    indexIVFPQ_stats.search_cycles += TOC;

    TIC;

    // 3rd level refinement
    size_t n_refine = 0;
#pragma omp parallel reduction(+ : n_refine)
    {
        // tmp buffers
        std::unique_ptr<float[]> residual_1(new float[2 * d]);
        float* residual_2 = residual_1.get() + d;
#pragma omp for
        for (idx_t i = 0; i < n; i++) {
            const float* xq = x + i * d;
            const idx_t* shortlist = coarse_labels.get() + k_coarse * i;
            float* heap_sim = distances + k * i;
            idx_t* heap_ids = labels + k * i;
            maxheap_heapify(k, heap_sim, heap_ids);

            for (int j = 0; j < k_coarse; j++) {
                idx_t sl = shortlist[j];

                if (sl == -1)
                    continue;

                int list_no = lo_listno(sl);
                int ofs = lo_offset(sl);

                assert(list_no >= 0 && list_no < nlist);
                assert(ofs >= 0 && ofs < invlists->list_size(list_no));

                // 1st level residual
                quantizer->compute_residual(xq, residual_1.get(), list_no);

                // 2nd level residual
                const uint8_t* l2code = invlists->get_single_code(list_no, ofs);

                pq.decode(l2code, residual_2);
                for (int l = 0; l < d; l++)
                    residual_2[l] = residual_1[l] - residual_2[l];

                // 3rd level residual's approximation
                idx_t id = invlists->get_single_id(list_no, ofs);
                assert(0 <= id && id < ntotal);
                refine_pq.decode(
                        &refine_codes[id * refine_pq.code_size],
                        residual_1.get());

                float dis = fvec_L2sqr(residual_1.get(), residual_2, d);

                if (dis < heap_sim[0]) {
                    idx_t id_or_pair = store_pairs ? sl : id;
                    maxheap_replace_top(k, heap_sim, heap_ids, dis, id_or_pair);
                }
                n_refine++;
            }
            maxheap_reorder(k, heap_sim, heap_ids);
        }
    }
    indexIVFPQ_stats.nrefine += n_refine;
    indexIVFPQ_stats.refine_cycles += TOC;
}

void IndexIVFPQR::reconstruct_from_offset(
        int64_t list_no,
        int64_t offset,
        float* recons) const {
    IndexIVFPQ::reconstruct_from_offset(list_no, offset, recons);

    idx_t id = invlists->get_single_id(list_no, offset);
    assert(0 <= id && id < ntotal);

    std::vector<float> r3(d);
    refine_pq.decode(&refine_codes[id * refine_pq.code_size], r3.data());
    for (int i = 0; i < d; ++i) {
        recons[i] += r3[i];
    }
}

void IndexIVFPQR::merge_from(Index& otherIndex, idx_t add_id) {
    IndexIVFPQR* other = dynamic_cast<IndexIVFPQR*>(&otherIndex);
    FAISS_THROW_IF_NOT(other);

    IndexIVF::merge_from(otherIndex, add_id);

    refine_codes.insert(
            refine_codes.end(),
            other->refine_codes.begin(),
            other->refine_codes.end());
    other->refine_codes.clear();
}

size_t IndexIVFPQR::remove_ids(const IDSelector& /*sel*/) {
    FAISS_THROW_MSG("not implemented");
    return 0;
}

} // namespace faiss
