import sys

sys.path.append('./')
import os

import faiss
import numpy as np
from arguments import IndexArguments, DataArguments, ModelArguments, TrainingArguments
from transformers import HfArgumentParser
from evaluate import validate, load_test_data

from LibVQ.base_index import FaissIndex

faiss.omp_set_num_threads(32)

if __name__ == '__main__':
    parser = HfArgumentParser((IndexArguments, DataArguments, ModelArguments, TrainingArguments))
    index_args, data_args, model_args, training_args = parser.parse_args_into_dataclasses()

    # Load embeddings
    emb_size = 768
    doc_embeddings = np.memmap(os.path.join(data_args.embeddings_dir, 'docs.memmap'),
                               dtype=np.float32, mode="r")
    doc_embeddings = doc_embeddings.reshape(-1, emb_size)
    query_embeddings = np.memmap(os.path.join(data_args.embeddings_dir, 'test-queries.memmap'),
                                 dtype=np.float32, mode="r")
    query_embeddings = query_embeddings.reshape(-1, emb_size)

    # Creat Faiss index
    index = FaissIndex(index_method=index_args.index_method,
                       emb_size=len(doc_embeddings[0]),
                       ivf_centers_num=index_args.ivf_centers_num,
                       subvector_num=index_args.subvector_num,
                       subvector_bits=index_args.subvector_bits,
                       dist_mode=index_args.dist_mode)

    print('Training the index with doc embeddings')
    # index.CPU_to_GPU(0)
    index.fit(doc_embeddings)
    index.add(doc_embeddings)
    # index.GPU_to_CPU()

    index_file = os.path.join(data_args.embeddings_dir,
                              f'{index_args.index_method}_ivf{index_args.ivf_centers_num}_pq{index_args.subvector_num}x{index_args.subvector_bits}.index')
    index.save_index(index_file)
    # index.load_index(index_file)

    # Test the performance
    scores, ann_items = index.search(query_embeddings, topk=100, nprobe=index_args.nprobe)
    test_questions, test_answers, collections = load_test_data(
        query_andwer_file='./data/NQ/raw_dataset/nq-test.qa.csv',
        collections_file='./data/NQ/dataset/collection.tsv')
    validate(ann_items, test_questions, test_answers, collections)
