/**
 * 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.
 */

// -*- c++ -*-

#include <faiss/index_io.h>

#include <faiss/impl/io.h>
#include <faiss/impl/io_macros.h>

#include <cstdio>
#include <cstdlib>

#include <sys/stat.h>
#include <sys/types.h>

#include <faiss/invlists/InvertedListsIOHook.h>

#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/io.h>
#include <faiss/impl/io_macros.h>
#include <faiss/utils/hamming.h>

#include <faiss/Index2Layer.h>
#include <faiss/IndexAdditiveQuantizer.h>
#include <faiss/IndexAdditiveQuantizerFastScan.h>
#include <faiss/IndexFlat.h>
#include <faiss/IndexHNSW.h>
#include <faiss/IndexIVF.h>
#include <faiss/IndexIVFAdditiveQuantizer.h>
#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/IndexIVFPQFastScan.h>
#include <faiss/IndexIVFPQR.h>
#include <faiss/IndexIVFSpectralHash.h>
#include <faiss/IndexLSH.h>
#include <faiss/IndexLattice.h>
#include <faiss/IndexNNDescent.h>
#include <faiss/IndexNSG.h>
#include <faiss/IndexPQ.h>
#include <faiss/IndexPQFastScan.h>
#include <faiss/IndexPreTransform.h>
#include <faiss/IndexRefine.h>
#include <faiss/IndexRowwiseMinMax.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/MetaIndexes.h>
#include <faiss/VectorTransform.h>

#include <faiss/IndexBinaryFlat.h>
#include <faiss/IndexBinaryFromFloat.h>
#include <faiss/IndexBinaryHNSW.h>
#include <faiss/IndexBinaryHash.h>
#include <faiss/IndexBinaryIVF.h>

/*************************************************************
 * The I/O format is the content of the class. For objects that are
 * inherited, like Index, a 4-character-code (fourcc) indicates which
 * child class this is an instance of.
 *
 * In this case, the fields of the parent class are written first,
 * then the ones for the child classes. Note that this requires
 * classes to be serialized to have a constructor without parameters,
 * so that the fields can be filled in later. The default constructor
 * should set reasonable defaults for all fields.
 *
 * The fourccs are assigned arbitrarily. When the class changed (added
 * or deprecated fields), the fourcc can be replaced. New code should
 * be able to read the old fourcc and fill in new classes.
 *
 * TODO: in this file, the read functions that encouter errors may
 * leak memory.
 **************************************************************/

namespace faiss {

/*************************************************************
 * Write
 **************************************************************/
static void write_index_header(const Index* idx, IOWriter* f) {
    WRITE1(idx->d);
    WRITE1(idx->ntotal);
    Index::idx_t dummy = 1 << 20;
    WRITE1(dummy);
    WRITE1(dummy);
    WRITE1(idx->is_trained);
    WRITE1(idx->metric_type);
    if (idx->metric_type > 1) {
        WRITE1(idx->metric_arg);
    }
}

void write_VectorTransform(const VectorTransform* vt, IOWriter* f) {
    if (const LinearTransform* lt = dynamic_cast<const LinearTransform*>(vt)) {
        if (dynamic_cast<const RandomRotationMatrix*>(lt)) {
            uint32_t h = fourcc("rrot");
            WRITE1(h);
        } else if (const PCAMatrix* pca = dynamic_cast<const PCAMatrix*>(lt)) {
            uint32_t h = fourcc("Pcam");
            WRITE1(h);
            WRITE1(pca->eigen_power);
            WRITE1(pca->epsilon);
            WRITE1(pca->random_rotation);
            WRITE1(pca->balanced_bins);
            WRITEVECTOR(pca->mean);
            WRITEVECTOR(pca->eigenvalues);
            WRITEVECTOR(pca->PCAMat);
        } else if (const ITQMatrix* itqm = dynamic_cast<const ITQMatrix*>(lt)) {
            uint32_t h = fourcc("Viqm");
            WRITE1(h);
            WRITE1(itqm->max_iter);
            WRITE1(itqm->seed);
        } else {
            // generic LinearTransform (includes OPQ)
            uint32_t h = fourcc("LTra");
            WRITE1(h);
        }
        WRITE1(lt->have_bias);
        WRITEVECTOR(lt->A);
        WRITEVECTOR(lt->b);
    } else if (
            const RemapDimensionsTransform* rdt =
                    dynamic_cast<const RemapDimensionsTransform*>(vt)) {
        uint32_t h = fourcc("RmDT");
        WRITE1(h);
        WRITEVECTOR(rdt->map);
    } else if (
            const NormalizationTransform* nt =
                    dynamic_cast<const NormalizationTransform*>(vt)) {
        uint32_t h = fourcc("VNrm");
        WRITE1(h);
        WRITE1(nt->norm);
    } else if (
            const CenteringTransform* ct =
                    dynamic_cast<const CenteringTransform*>(vt)) {
        uint32_t h = fourcc("VCnt");
        WRITE1(h);
        WRITEVECTOR(ct->mean);
    } else if (
            const ITQTransform* itqt = dynamic_cast<const ITQTransform*>(vt)) {
        uint32_t h = fourcc("Viqt");
        WRITE1(h);
        WRITEVECTOR(itqt->mean);
        WRITE1(itqt->do_pca);
        write_VectorTransform(&itqt->itq, f);
        write_VectorTransform(&itqt->pca_then_itq, f);
    } else {
        FAISS_THROW_MSG("cannot serialize this");
    }
    // common fields
    WRITE1(vt->d_in);
    WRITE1(vt->d_out);
    WRITE1(vt->is_trained);
}

void write_ProductQuantizer(const ProductQuantizer* pq, IOWriter* f) {
    WRITE1(pq->d);
    WRITE1(pq->M);
    WRITE1(pq->nbits);
    WRITEVECTOR(pq->centroids);
}

static void write_AdditiveQuantizer(const AdditiveQuantizer* aq, IOWriter* f) {
    WRITE1(aq->d);
    WRITE1(aq->M);
    WRITEVECTOR(aq->nbits);
    WRITE1(aq->is_trained);
    WRITEVECTOR(aq->codebooks);
    WRITE1(aq->search_type);
    WRITE1(aq->norm_min);
    WRITE1(aq->norm_max);
    if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8 ||
        aq->search_type == AdditiveQuantizer::ST_norm_cqint4 ||
        aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 ||
        aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) {
        WRITEXBVECTOR(aq->qnorm.codes);
    }

    if (aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 ||
        aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) {
        WRITEVECTOR(aq->norm_tabs);
    }
}

static void write_ResidualQuantizer(const ResidualQuantizer* rq, IOWriter* f) {
    write_AdditiveQuantizer(rq, f);
    WRITE1(rq->train_type);
    WRITE1(rq->max_beam_size);
}

static void write_LocalSearchQuantizer(
        const LocalSearchQuantizer* lsq,
        IOWriter* f) {
    write_AdditiveQuantizer(lsq, f);
    WRITE1(lsq->K);
    WRITE1(lsq->train_iters);
    WRITE1(lsq->encode_ils_iters);
    WRITE1(lsq->train_ils_iters);
    WRITE1(lsq->icm_iters);
    WRITE1(lsq->p);
    WRITE1(lsq->lambd);
    WRITE1(lsq->chunk_size);
    WRITE1(lsq->random_seed);
    WRITE1(lsq->nperts);
    WRITE1(lsq->update_codebooks_with_double);
}

static void write_ProductAdditiveQuantizer(
        const ProductAdditiveQuantizer* paq,
        IOWriter* f) {
    write_AdditiveQuantizer(paq, f);
    WRITE1(paq->nsplits);
}

static void write_ProductResidualQuantizer(
        const ProductResidualQuantizer* prq,
        IOWriter* f) {
    write_ProductAdditiveQuantizer(prq, f);
    for (const auto aq : prq->quantizers) {
        auto rq = dynamic_cast<const ResidualQuantizer*>(aq);
        write_ResidualQuantizer(rq, f);
    }
}

static void write_ProductLocalSearchQuantizer(
        const ProductLocalSearchQuantizer* plsq,
        IOWriter* f) {
    write_ProductAdditiveQuantizer(plsq, f);
    for (const auto aq : plsq->quantizers) {
        auto lsq = dynamic_cast<const LocalSearchQuantizer*>(aq);
        write_LocalSearchQuantizer(lsq, f);
    }
}

static void write_ScalarQuantizer(const ScalarQuantizer* ivsc, IOWriter* f) {
    WRITE1(ivsc->qtype);
    WRITE1(ivsc->rangestat);
    WRITE1(ivsc->rangestat_arg);
    WRITE1(ivsc->d);
    WRITE1(ivsc->code_size);
    WRITEVECTOR(ivsc->trained);
}

void write_InvertedLists(const InvertedLists* ils, IOWriter* f) {
    if (ils == nullptr) {
        uint32_t h = fourcc("il00");
        WRITE1(h);
    } else if (
            const auto& ails = dynamic_cast<const ArrayInvertedLists*>(ils)) {
        uint32_t h = fourcc("ilar");
        WRITE1(h);
        WRITE1(ails->nlist);
        WRITE1(ails->code_size);
        // here we store either as a full or a sparse data buffer
        size_t n_non0 = 0;
        for (size_t i = 0; i < ails->nlist; i++) {
            if (ails->ids[i].size() > 0)
                n_non0++;
        }
        if (n_non0 > ails->nlist / 2) {
            uint32_t list_type = fourcc("full");
            WRITE1(list_type);
            std::vector<size_t> sizes;
            for (size_t i = 0; i < ails->nlist; i++) {
                sizes.push_back(ails->ids[i].size());
            }
            WRITEVECTOR(sizes);
        } else {
            int list_type = fourcc("sprs"); // sparse
            WRITE1(list_type);
            std::vector<size_t> sizes;
            for (size_t i = 0; i < ails->nlist; i++) {
                size_t n = ails->ids[i].size();
                if (n > 0) {
                    sizes.push_back(i);
                    sizes.push_back(n);
                }
            }
            WRITEVECTOR(sizes);
        }
        // make a single contiguous data buffer (useful for mmapping)
        for (size_t i = 0; i < ails->nlist; i++) {
            size_t n = ails->ids[i].size();
            if (n > 0) {
                WRITEANDCHECK(ails->codes[i].data(), n * ails->code_size);
                WRITEANDCHECK(ails->ids[i].data(), n);
            }
        }

    } else {
        InvertedListsIOHook::lookup_classname(typeid(*ils).name())
                ->write(ils, f);
    }
}

void write_ProductQuantizer(const ProductQuantizer* pq, const char* fname) {
    FileIOWriter writer(fname);
    write_ProductQuantizer(pq, &writer);
}

static void write_HNSW(const HNSW* hnsw, IOWriter* f) {
    WRITEVECTOR(hnsw->assign_probas);
    WRITEVECTOR(hnsw->cum_nneighbor_per_level);
    WRITEVECTOR(hnsw->levels);
    WRITEVECTOR(hnsw->offsets);
    WRITEVECTOR(hnsw->neighbors);

    WRITE1(hnsw->entry_point);
    WRITE1(hnsw->max_level);
    WRITE1(hnsw->efConstruction);
    WRITE1(hnsw->efSearch);
    WRITE1(hnsw->upper_beam);
}

static void write_NSG(const NSG* nsg, IOWriter* f) {
    WRITE1(nsg->ntotal);
    WRITE1(nsg->R);
    WRITE1(nsg->L);
    WRITE1(nsg->C);
    WRITE1(nsg->search_L);
    WRITE1(nsg->enterpoint);
    WRITE1(nsg->is_built);

    if (!nsg->is_built) {
        return;
    }

    constexpr int EMPTY_ID = -1;
    auto& graph = nsg->final_graph;
    int K = graph->K;
    int N = graph->N;
    FAISS_THROW_IF_NOT(N == nsg->ntotal);
    FAISS_THROW_IF_NOT(K == nsg->R);
    FAISS_THROW_IF_NOT(true == graph->own_fields);

    int size = 0;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < K; j++) {
            int id = graph->at(i, j);
            if (id != EMPTY_ID) {
                WRITE1(id);
                size += 1;
            } else {
                break;
            }
        }
        WRITE1(EMPTY_ID);
    }
}

static void write_NNDescent(const NNDescent* nnd, IOWriter* f) {
    WRITE1(nnd->ntotal);
    WRITE1(nnd->d);
    WRITE1(nnd->K);
    WRITE1(nnd->S);
    WRITE1(nnd->R);
    WRITE1(nnd->L);
    WRITE1(nnd->iter);
    WRITE1(nnd->search_L);
    WRITE1(nnd->random_seed);
    WRITE1(nnd->has_built);

    WRITEVECTOR(nnd->final_graph);
}

static void write_direct_map(const DirectMap* dm, IOWriter* f) {
    char maintain_direct_map =
            (char)dm->type; // for backwards compatibility with bool
    WRITE1(maintain_direct_map);
    WRITEVECTOR(dm->array);
    if (dm->type == DirectMap::Hashtable) {
        using idx_t = Index::idx_t;
        std::vector<std::pair<idx_t, idx_t>> v;
        const std::unordered_map<idx_t, idx_t>& map = dm->hashtable;
        v.resize(map.size());
        std::copy(map.begin(), map.end(), v.begin());
        WRITEVECTOR(v);
    }
}

static void write_ivf_header(const IndexIVF* ivf, IOWriter* f) {
    write_index_header(ivf, f);
    WRITE1(ivf->nlist);
    WRITE1(ivf->nprobe);
    write_index(ivf->quantizer, f);
    write_direct_map(&ivf->direct_map, f);
}

void write_index(const Index* idx, IOWriter* f) {
    if (const IndexFlat* idxf = dynamic_cast<const IndexFlat*>(idx)) {
        uint32_t h =
                fourcc(idxf->metric_type == METRIC_INNER_PRODUCT ? "IxFI"
                               : idxf->metric_type == METRIC_L2  ? "IxF2"
                                                                 : "IxFl");
        WRITE1(h);
        write_index_header(idx, f);
        WRITEXBVECTOR(idxf->codes);
    } else if (const IndexLSH* idxl = dynamic_cast<const IndexLSH*>(idx)) {
        uint32_t h = fourcc("IxHe");
        WRITE1(h);
        write_index_header(idx, f);
        WRITE1(idxl->nbits);
        WRITE1(idxl->rotate_data);
        WRITE1(idxl->train_thresholds);
        WRITEVECTOR(idxl->thresholds);
        int code_size_i = idxl->code_size;
        WRITE1(code_size_i);
        write_VectorTransform(&idxl->rrot, f);
        WRITEVECTOR(idxl->codes);
    } else if (const IndexPQ* idxp = dynamic_cast<const IndexPQ*>(idx)) {
        uint32_t h = fourcc("IxPq");
        WRITE1(h);
        write_index_header(idx, f);
        write_ProductQuantizer(&idxp->pq, f);
        WRITEVECTOR(idxp->codes);
        // search params -- maybe not useful to store?
        WRITE1(idxp->search_type);
        WRITE1(idxp->encode_signs);
        WRITE1(idxp->polysemous_ht);
    } else if (
            const IndexResidualQuantizer* idxr =
                    dynamic_cast<const IndexResidualQuantizer*>(idx)) {
        uint32_t h = fourcc("IxRq");
        WRITE1(h);
        write_index_header(idx, f);
        write_ResidualQuantizer(&idxr->rq, f);
        WRITE1(idxr->code_size);
        WRITEVECTOR(idxr->codes);
    } else if (
            auto* idxr = dynamic_cast<const IndexLocalSearchQuantizer*>(idx)) {
        uint32_t h = fourcc("IxLS");
        WRITE1(h);
        write_index_header(idx, f);
        write_LocalSearchQuantizer(&idxr->lsq, f);
        WRITE1(idxr->code_size);
        WRITEVECTOR(idxr->codes);
    } else if (
            const IndexProductResidualQuantizer* idxpr =
                    dynamic_cast<const IndexProductResidualQuantizer*>(idx)) {
        uint32_t h = fourcc("IxPR");
        WRITE1(h);
        write_index_header(idx, f);
        write_ProductResidualQuantizer(&idxpr->prq, f);
        WRITE1(idxpr->code_size);
        WRITEVECTOR(idxpr->codes);
    } else if (
            const IndexProductLocalSearchQuantizer* idxpl =
                    dynamic_cast<const IndexProductLocalSearchQuantizer*>(
                            idx)) {
        uint32_t h = fourcc("IxPL");
        WRITE1(h);
        write_index_header(idx, f);
        write_ProductLocalSearchQuantizer(&idxpl->plsq, f);
        WRITE1(idxpl->code_size);
        WRITEVECTOR(idxpl->codes);
    } else if (
            auto* idxaqfs =
                    dynamic_cast<const IndexAdditiveQuantizerFastScan*>(idx)) {
        auto idxlsqfs =
                dynamic_cast<const IndexLocalSearchQuantizerFastScan*>(idx);
        auto idxrqfs = dynamic_cast<const IndexResidualQuantizerFastScan*>(idx);
        auto idxplsqfs =
                dynamic_cast<const IndexProductLocalSearchQuantizerFastScan*>(
                        idx);
        auto idxprqfs =
                dynamic_cast<const IndexProductResidualQuantizerFastScan*>(idx);
        FAISS_THROW_IF_NOT(idxlsqfs || idxrqfs || idxplsqfs || idxprqfs);

        if (idxlsqfs) {
            uint32_t h = fourcc("ILfs");
            WRITE1(h);
        } else if (idxrqfs) {
            uint32_t h = fourcc("IRfs");
            WRITE1(h);
        } else if (idxplsqfs) {
            uint32_t h = fourcc("IPLf");
            WRITE1(h);
        } else if (idxprqfs) {
            uint32_t h = fourcc("IPRf");
            WRITE1(h);
        }

        write_index_header(idxaqfs, f);

        if (idxlsqfs) {
            write_LocalSearchQuantizer(&idxlsqfs->lsq, f);
        } else if (idxrqfs) {
            write_ResidualQuantizer(&idxrqfs->rq, f);
        } else if (idxplsqfs) {
            write_ProductLocalSearchQuantizer(&idxplsqfs->plsq, f);
        } else if (idxprqfs) {
            write_ProductResidualQuantizer(&idxprqfs->prq, f);
        }
        WRITE1(idxaqfs->implem);
        WRITE1(idxaqfs->bbs);
        WRITE1(idxaqfs->qbs);

        WRITE1(idxaqfs->M);
        WRITE1(idxaqfs->nbits);
        WRITE1(idxaqfs->ksub);
        WRITE1(idxaqfs->code_size);
        WRITE1(idxaqfs->ntotal2);
        WRITE1(idxaqfs->M2);

        WRITE1(idxaqfs->rescale_norm);
        WRITE1(idxaqfs->norm_scale);
        WRITE1(idxaqfs->max_train_points);

        WRITEVECTOR(idxaqfs->codes);
    } else if (
            auto* ivaqfs =
                    dynamic_cast<const IndexIVFAdditiveQuantizerFastScan*>(
                            idx)) {
        auto ivlsqfs =
                dynamic_cast<const IndexIVFLocalSearchQuantizerFastScan*>(idx);
        auto ivrqfs =
                dynamic_cast<const IndexIVFResidualQuantizerFastScan*>(idx);
        auto ivplsqfs = dynamic_cast<
                const IndexIVFProductLocalSearchQuantizerFastScan*>(idx);
        auto ivprqfs =
                dynamic_cast<const IndexIVFProductResidualQuantizerFastScan*>(
                        idx);
        FAISS_THROW_IF_NOT(ivlsqfs || ivrqfs || ivplsqfs || ivprqfs);

        if (ivlsqfs) {
            uint32_t h = fourcc("IVLf");
            WRITE1(h);
        } else if (ivrqfs) {
            uint32_t h = fourcc("IVRf");
            WRITE1(h);
        } else if (ivplsqfs) {
            uint32_t h = fourcc("NPLf"); // N means IV ...
            WRITE1(h);
        } else {
            uint32_t h = fourcc("NPRf");
            WRITE1(h);
        }

        write_ivf_header(ivaqfs, f);

        if (ivlsqfs) {
            write_LocalSearchQuantizer(&ivlsqfs->lsq, f);
        } else if (ivrqfs) {
            write_ResidualQuantizer(&ivrqfs->rq, f);
        } else if (ivplsqfs) {
            write_ProductLocalSearchQuantizer(&ivplsqfs->plsq, f);
        } else {
            write_ProductResidualQuantizer(&ivprqfs->prq, f);
        }

        WRITE1(ivaqfs->by_residual);
        WRITE1(ivaqfs->implem);
        WRITE1(ivaqfs->bbs);
        WRITE1(ivaqfs->qbs);

        WRITE1(ivaqfs->M);
        WRITE1(ivaqfs->nbits);
        WRITE1(ivaqfs->ksub);
        WRITE1(ivaqfs->code_size);
        WRITE1(ivaqfs->qbs2);
        WRITE1(ivaqfs->M2);

        WRITE1(ivaqfs->rescale_norm);
        WRITE1(ivaqfs->norm_scale);
        WRITE1(ivaqfs->max_train_points);

        write_InvertedLists(ivaqfs->invlists, f);
    } else if (
            const ResidualCoarseQuantizer* idxr =
                    dynamic_cast<const ResidualCoarseQuantizer*>(idx)) {
        uint32_t h = fourcc("ImRQ");
        WRITE1(h);
        write_index_header(idx, f);
        write_ResidualQuantizer(&idxr->rq, f);
        WRITE1(idxr->beam_factor);
    } else if (
            const Index2Layer* idxp = dynamic_cast<const Index2Layer*>(idx)) {
        uint32_t h = fourcc("Ix2L");
        WRITE1(h);
        write_index_header(idx, f);
        write_index(idxp->q1.quantizer, f);
        WRITE1(idxp->q1.nlist);
        WRITE1(idxp->q1.quantizer_trains_alone);
        write_ProductQuantizer(&idxp->pq, f);
        WRITE1(idxp->code_size_1);
        WRITE1(idxp->code_size_2);
        WRITE1(idxp->code_size);
        WRITEVECTOR(idxp->codes);
    } else if (
            const IndexScalarQuantizer* idxs =
                    dynamic_cast<const IndexScalarQuantizer*>(idx)) {
        uint32_t h = fourcc("IxSQ");
        WRITE1(h);
        write_index_header(idx, f);
        write_ScalarQuantizer(&idxs->sq, f);
        WRITEVECTOR(idxs->codes);
    } else if (
            const IndexLattice* idxl = dynamic_cast<const IndexLattice*>(idx)) {
        uint32_t h = fourcc("IxLa");
        WRITE1(h);
        WRITE1(idxl->d);
        WRITE1(idxl->nsq);
        WRITE1(idxl->scale_nbit);
        WRITE1(idxl->zn_sphere_codec.r2);
        write_index_header(idx, f);
        WRITEVECTOR(idxl->trained);
    } else if (
            const IndexIVFFlatDedup* ivfl =
                    dynamic_cast<const IndexIVFFlatDedup*>(idx)) {
        uint32_t h = fourcc("IwFd");
        WRITE1(h);
        write_ivf_header(ivfl, f);
        {
            std::vector<Index::idx_t> tab(2 * ivfl->instances.size());
            long i = 0;
            for (auto it = ivfl->instances.begin(); it != ivfl->instances.end();
                 ++it) {
                tab[i++] = it->first;
                tab[i++] = it->second;
            }
            WRITEVECTOR(tab);
        }
        write_InvertedLists(ivfl->invlists, f);
    } else if (
            const IndexIVFFlat* ivfl = dynamic_cast<const IndexIVFFlat*>(idx)) {
        uint32_t h = fourcc("IwFl");
        WRITE1(h);
        write_ivf_header(ivfl, f);
        write_InvertedLists(ivfl->invlists, f);
    } else if (
            const IndexIVFScalarQuantizer* ivsc =
                    dynamic_cast<const IndexIVFScalarQuantizer*>(idx)) {
        uint32_t h = fourcc("IwSq");
        WRITE1(h);
        write_ivf_header(ivsc, f);
        write_ScalarQuantizer(&ivsc->sq, f);
        WRITE1(ivsc->code_size);
        WRITE1(ivsc->by_residual);
        write_InvertedLists(ivsc->invlists, f);
    } else if (auto iva = dynamic_cast<const IndexIVFAdditiveQuantizer*>(idx)) {
        bool is_LSQ = dynamic_cast<const IndexIVFLocalSearchQuantizer*>(iva);
        bool is_RQ = dynamic_cast<const IndexIVFResidualQuantizer*>(iva);
        bool is_PLSQ =
                dynamic_cast<const IndexIVFProductLocalSearchQuantizer*>(iva);
        uint32_t h;
        if (is_LSQ) {
            h = fourcc("IwLS");
        } else if (is_RQ) {
            h = fourcc("IwRQ");
        } else if (is_PLSQ) {
            h = fourcc("IwPL");
        } else {
            h = fourcc("IwPR");
        }

        WRITE1(h);
        write_ivf_header(iva, f);
        WRITE1(iva->code_size);
        if (is_LSQ) {
            write_LocalSearchQuantizer((LocalSearchQuantizer*)iva->aq, f);
        } else if (is_RQ) {
            write_ResidualQuantizer((ResidualQuantizer*)iva->aq, f);
        } else if (is_PLSQ) {
            write_ProductLocalSearchQuantizer(
                    (ProductLocalSearchQuantizer*)iva->aq, f);
        } else {
            write_ProductResidualQuantizer(
                    (ProductResidualQuantizer*)iva->aq, f);
        }
        WRITE1(iva->by_residual);
        WRITE1(iva->use_precomputed_table);
        write_InvertedLists(iva->invlists, f);
    } else if (
            const IndexIVFSpectralHash* ivsp =
                    dynamic_cast<const IndexIVFSpectralHash*>(idx)) {
        uint32_t h = fourcc("IwSh");
        WRITE1(h);
        write_ivf_header(ivsp, f);
        write_VectorTransform(ivsp->vt, f);
        WRITE1(ivsp->nbit);
        WRITE1(ivsp->period);
        WRITE1(ivsp->threshold_type);
        WRITEVECTOR(ivsp->trained);
        write_InvertedLists(ivsp->invlists, f);
    } else if (const IndexIVFPQ* ivpq = dynamic_cast<const IndexIVFPQ*>(idx)) {
        const IndexIVFPQR* ivfpqr = dynamic_cast<const IndexIVFPQR*>(idx);

        uint32_t h = fourcc(ivfpqr ? "IwQR" : "IwPQ");
        WRITE1(h);
        write_ivf_header(ivpq, f);
        WRITE1(ivpq->by_residual);
        WRITE1(ivpq->code_size);
        write_ProductQuantizer(&ivpq->pq, f);
        write_InvertedLists(ivpq->invlists, f);
        if (ivfpqr) {
            write_ProductQuantizer(&ivfpqr->refine_pq, f);
            WRITEVECTOR(ivfpqr->refine_codes);
            WRITE1(ivfpqr->k_factor);
        }

    } else if (
            const IndexPreTransform* ixpt =
                    dynamic_cast<const IndexPreTransform*>(idx)) {
        uint32_t h = fourcc("IxPT");
        WRITE1(h);
        write_index_header(ixpt, f);
        int nt = ixpt->chain.size();
        WRITE1(nt);
        for (int i = 0; i < nt; i++)
            write_VectorTransform(ixpt->chain[i], f);
        write_index(ixpt->index, f);
    } else if (
            const MultiIndexQuantizer* imiq =
                    dynamic_cast<const MultiIndexQuantizer*>(idx)) {
        uint32_t h = fourcc("Imiq");
        WRITE1(h);
        write_index_header(imiq, f);
        write_ProductQuantizer(&imiq->pq, f);
    } else if (
            const IndexRefine* idxrf = dynamic_cast<const IndexRefine*>(idx)) {
        uint32_t h = fourcc("IxRF");
        WRITE1(h);
        write_index_header(idxrf, f);
        write_index(idxrf->base_index, f);
        write_index(idxrf->refine_index, f);
        WRITE1(idxrf->k_factor);
    } else if (
            const IndexIDMap* idxmap = dynamic_cast<const IndexIDMap*>(idx)) {
        uint32_t h = dynamic_cast<const IndexIDMap2*>(idx) ? fourcc("IxM2")
                                                           : fourcc("IxMp");
        // no need to store additional info for IndexIDMap2
        WRITE1(h);
        write_index_header(idxmap, f);
        write_index(idxmap->index, f);
        WRITEVECTOR(idxmap->id_map);
    } else if (const IndexHNSW* idxhnsw = dynamic_cast<const IndexHNSW*>(idx)) {
        uint32_t h = dynamic_cast<const IndexHNSWFlat*>(idx) ? fourcc("IHNf")
                : dynamic_cast<const IndexHNSWPQ*>(idx)      ? fourcc("IHNp")
                : dynamic_cast<const IndexHNSWSQ*>(idx)      ? fourcc("IHNs")
                : dynamic_cast<const IndexHNSW2Level*>(idx)  ? fourcc("IHN2")
                                                             : 0;
        FAISS_THROW_IF_NOT(h != 0);
        WRITE1(h);
        write_index_header(idxhnsw, f);
        write_HNSW(&idxhnsw->hnsw, f);
        write_index(idxhnsw->storage, f);
    } else if (const IndexNSG* idxnsg = dynamic_cast<const IndexNSG*>(idx)) {
        uint32_t h = dynamic_cast<const IndexNSGFlat*>(idx) ? fourcc("INSf")
                : dynamic_cast<const IndexNSGPQ*>(idx)      ? fourcc("INSp")
                : dynamic_cast<const IndexNSGSQ*>(idx)      ? fourcc("INSs")
                                                            : 0;
        FAISS_THROW_IF_NOT(h != 0);
        WRITE1(h);
        write_index_header(idxnsg, f);
        WRITE1(idxnsg->GK);
        WRITE1(idxnsg->build_type);
        WRITE1(idxnsg->nndescent_S);
        WRITE1(idxnsg->nndescent_R);
        WRITE1(idxnsg->nndescent_L);
        WRITE1(idxnsg->nndescent_iter);
        write_NSG(&idxnsg->nsg, f);
        write_index(idxnsg->storage, f);
    } else if (
            const IndexNNDescent* idxnnd =
                    dynamic_cast<const IndexNNDescent*>(idx)) {
        auto idxnndflat = dynamic_cast<const IndexNNDescentFlat*>(idx);
        FAISS_THROW_IF_NOT(idxnndflat != nullptr);
        uint32_t h = fourcc("INNf");
        FAISS_THROW_IF_NOT(h != 0);
        WRITE1(h);
        write_index_header(idxnnd, f);
        write_NNDescent(&idxnnd->nndescent, f);
        write_index(idxnnd->storage, f);
    } else if (
            const IndexPQFastScan* idxpqfs =
                    dynamic_cast<const IndexPQFastScan*>(idx)) {
        uint32_t h = fourcc("IPfs");
        WRITE1(h);
        write_index_header(idxpqfs, f);
        write_ProductQuantizer(&idxpqfs->pq, f);
        WRITE1(idxpqfs->implem);
        WRITE1(idxpqfs->bbs);
        WRITE1(idxpqfs->qbs);
        WRITE1(idxpqfs->ntotal2);
        WRITE1(idxpqfs->M2);
        WRITEVECTOR(idxpqfs->codes);
    } else if (
            const IndexIVFPQFastScan* ivpq =
                    dynamic_cast<const IndexIVFPQFastScan*>(idx)) {
        uint32_t h = fourcc("IwPf");
        WRITE1(h);
        write_ivf_header(ivpq, f);
        WRITE1(ivpq->by_residual);
        WRITE1(ivpq->code_size);
        WRITE1(ivpq->bbs);
        WRITE1(ivpq->M2);
        WRITE1(ivpq->implem);
        WRITE1(ivpq->qbs2);
        write_ProductQuantizer(&ivpq->pq, f);
        write_InvertedLists(ivpq->invlists, f);
    } else if (
            const IndexRowwiseMinMax* imm =
                    dynamic_cast<const IndexRowwiseMinMax*>(idx)) {
        // IndexRowwiseMinmaxFloat
        uint32_t h = fourcc("IRMf");
        WRITE1(h);
        write_index_header(imm, f);
        write_index(imm->index, f);
    } else if (
            const IndexRowwiseMinMaxFP16* imm =
                    dynamic_cast<const IndexRowwiseMinMaxFP16*>(idx)) {
        // IndexRowwiseMinmaxHalf
        uint32_t h = fourcc("IRMh");
        WRITE1(h);
        write_index_header(imm, f);
        write_index(imm->index, f);
    } else {
        FAISS_THROW_MSG("don't know how to serialize this type of index");
    }
}

void write_index(const Index* idx, FILE* f) {
    FileIOWriter writer(f);
    write_index(idx, &writer);
}

void write_index(const Index* idx, const char* fname) {
    FileIOWriter writer(fname);
    write_index(idx, &writer);
}

void write_VectorTransform(const VectorTransform* vt, const char* fname) {
    FileIOWriter writer(fname);
    write_VectorTransform(vt, &writer);
}

/*************************************************************
 * Write binary indexes
 **************************************************************/

static void write_index_binary_header(const IndexBinary* idx, IOWriter* f) {
    WRITE1(idx->d);
    WRITE1(idx->code_size);
    WRITE1(idx->ntotal);
    WRITE1(idx->is_trained);
    WRITE1(idx->metric_type);
}

static void write_binary_ivf_header(const IndexBinaryIVF* ivf, IOWriter* f) {
    write_index_binary_header(ivf, f);
    WRITE1(ivf->nlist);
    WRITE1(ivf->nprobe);
    write_index_binary(ivf->quantizer, f);
    write_direct_map(&ivf->direct_map, f);
}

static void write_binary_hash_invlists(
        const IndexBinaryHash::InvertedListMap& invlists,
        int b,
        IOWriter* f) {
    size_t sz = invlists.size();
    WRITE1(sz);
    size_t maxil = 0;
    for (auto it = invlists.begin(); it != invlists.end(); ++it) {
        if (it->second.ids.size() > maxil) {
            maxil = it->second.ids.size();
        }
    }
    int il_nbit = 0;
    while (maxil >= ((uint64_t)1 << il_nbit)) {
        il_nbit++;
    }
    WRITE1(il_nbit);

    // first write sizes then data, may be useful if we want to
    // memmap it at some point

    // buffer for bitstrings
    std::vector<uint8_t> buf(((b + il_nbit) * sz + 7) / 8);
    BitstringWriter wr(buf.data(), buf.size());
    for (auto it = invlists.begin(); it != invlists.end(); ++it) {
        wr.write(it->first, b);
        wr.write(it->second.ids.size(), il_nbit);
    }
    WRITEVECTOR(buf);

    for (auto it = invlists.begin(); it != invlists.end(); ++it) {
        WRITEVECTOR(it->second.ids);
        WRITEVECTOR(it->second.vecs);
    }
}

static void write_binary_multi_hash_map(
        const IndexBinaryMultiHash::Map& map,
        int b,
        size_t ntotal,
        IOWriter* f) {
    int id_bits = 0;
    while ((ntotal > ((Index::idx_t)1 << id_bits))) {
        id_bits++;
    }
    WRITE1(id_bits);
    size_t sz = map.size();
    WRITE1(sz);
    size_t nbit = (b + id_bits) * sz + ntotal * id_bits;
    std::vector<uint8_t> buf((nbit + 7) / 8);
    BitstringWriter wr(buf.data(), buf.size());
    for (auto it = map.begin(); it != map.end(); ++it) {
        wr.write(it->first, b);
        wr.write(it->second.size(), id_bits);
        for (auto id : it->second) {
            wr.write(id, id_bits);
        }
    }
    WRITEVECTOR(buf);
}

void write_index_binary(const IndexBinary* idx, IOWriter* f) {
    if (const IndexBinaryFlat* idxf =
                dynamic_cast<const IndexBinaryFlat*>(idx)) {
        uint32_t h = fourcc("IBxF");
        WRITE1(h);
        write_index_binary_header(idx, f);
        WRITEVECTOR(idxf->xb);
    } else if (
            const IndexBinaryIVF* ivf =
                    dynamic_cast<const IndexBinaryIVF*>(idx)) {
        uint32_t h = fourcc("IBwF");
        WRITE1(h);
        write_binary_ivf_header(ivf, f);
        write_InvertedLists(ivf->invlists, f);
    } else if (
            const IndexBinaryFromFloat* idxff =
                    dynamic_cast<const IndexBinaryFromFloat*>(idx)) {
        uint32_t h = fourcc("IBFf");
        WRITE1(h);
        write_index_binary_header(idxff, f);
        write_index(idxff->index, f);
    } else if (
            const IndexBinaryHNSW* idxhnsw =
                    dynamic_cast<const IndexBinaryHNSW*>(idx)) {
        uint32_t h = fourcc("IBHf");
        WRITE1(h);
        write_index_binary_header(idxhnsw, f);
        write_HNSW(&idxhnsw->hnsw, f);
        write_index_binary(idxhnsw->storage, f);
    } else if (
            const IndexBinaryIDMap* idxmap =
                    dynamic_cast<const IndexBinaryIDMap*>(idx)) {
        uint32_t h = dynamic_cast<const IndexBinaryIDMap2*>(idx)
                ? fourcc("IBM2")
                : fourcc("IBMp");
        // no need to store additional info for IndexIDMap2
        WRITE1(h);
        write_index_binary_header(idxmap, f);
        write_index_binary(idxmap->index, f);
        WRITEVECTOR(idxmap->id_map);
    } else if (
            const IndexBinaryHash* idxh =
                    dynamic_cast<const IndexBinaryHash*>(idx)) {
        uint32_t h = fourcc("IBHh");
        WRITE1(h);
        write_index_binary_header(idxh, f);
        WRITE1(idxh->b);
        WRITE1(idxh->nflip);
        write_binary_hash_invlists(idxh->invlists, idxh->b, f);
    } else if (
            const IndexBinaryMultiHash* idxmh =
                    dynamic_cast<const IndexBinaryMultiHash*>(idx)) {
        uint32_t h = fourcc("IBHm");
        WRITE1(h);
        write_index_binary_header(idxmh, f);
        write_index_binary(idxmh->storage, f);
        WRITE1(idxmh->b);
        WRITE1(idxmh->nhash);
        WRITE1(idxmh->nflip);
        for (int i = 0; i < idxmh->nhash; i++) {
            write_binary_multi_hash_map(
                    idxmh->maps[i], idxmh->b, idxmh->ntotal, f);
        }
    } else {
        FAISS_THROW_MSG("don't know how to serialize this type of index");
    }
}

void write_index_binary(const IndexBinary* idx, FILE* f) {
    FileIOWriter writer(f);
    write_index_binary(idx, &writer);
}

void write_index_binary(const IndexBinary* idx, const char* fname) {
    FileIOWriter writer(fname);
    write_index_binary(idx, &writer);
}

} // namespace faiss
