from datasets import load_dataset
from transformers import DPRContextEncoderTokenizer, DPRContextEncoder
from general_utils import embed_passages, embed_passages_haystack
import faiss
import argparse
import os
from haystack.nodes import DensePassageRetriever
from haystack.document_stores import InMemoryDocumentStore


os.environ["OMP_NUM_THREADS"] = "8"


def create_faiss_index(args):
    minchars = 200
    dims = 128

    dpr = DensePassageRetriever(
        document_store=InMemoryDocumentStore(),
        query_embedding_model="IIC/dpr-spanish-question_encoder-allqa-base",
        passage_embedding_model="IIC/dpr-spanish-question_encoder-allqa-base",
        max_seq_len_query=64,
        max_seq_len_passage=256,
        batch_size=512,
    )

    dataset = load_dataset(
        "IIC/spanish_biomedical_crawled_corpus", split="train"
    )

    dataset = dataset.filter(lambda example: len(example["text"]) > minchars)

    def embed_passages_retrieval(examples):
        return embed_passages_haystack(dpr, examples)

    dataset = dataset.map(embed_passages_retrieval, batched=True, batch_size=8192)

    dataset.add_faiss_index(
        column="embeddings",
        string_factory="OPQ64_128,IVF4898,PQ64x4fsr",
        train_size=len(dataset),
    )
    dataset.save_faiss_index("embeddings", args.index_file_name)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Creates Faiss Wikipedia index file")

    parser.add_argument(
        "--ctx_encoder_name",
        default="IIC/dpr-spanish-passage_encoder-squades-base",
        help="Encoding model to use for passage encoding",
    )

    parser.add_argument(
        "--index_file_name",
        default="dpr_index_bio_splitted.faiss",
        help="Faiss index file with passage embeddings",
    )
    parser.add_argument(
        "--device", default="cuda:0", help="The device to index data on."
    )

    main_args, _ = parser.parse_known_args()
    create_faiss_index(main_args)