import json
from datasets import load_dataset
from tqdm import tqdm


names = ["MCINext/FEVER_FA_test_top_250_only_w_correct-v2", "MCINext/fiqa-fa-v2", "MCINext/HotpotQA_FA_test_top_250_only_w_correct-v2",
"MCINext/MSMARCO_FA_test_top_250_only_w_correct-v2", "MCINext/NQ_FA_test_top_250_only_w_correct-v2", "MCINext/quora-fa-v2", "MCINext/scifact-fa-v2",
"MCINext/synthetic-persian-chatbot-rag-faq-retrieval", "MCINext/synthetic-persian-qa-retrieval", "MCINext/trec-covid-fa-v2"]

names = names[3:4]
for name in tqdm(names):
    print(f"loading {name}")
    dataset_qrel = load_dataset(name)["test"]
    dataset_corpus_list = load_dataset(name,data_files="corpus/corpus.jsonl")["train"]
    dataset_corpus = {}
    for data in dataset_corpus_list:
        dataset_corpus[data["_id"]] = data["text"]

    dataset_queries_list = load_dataset(name,data_files="queries/queries.jsonl")["train"]
    dataset_queries = {}
    for data in dataset_queries_list:
        dataset_queries[data["_id"]] = data["text"]


    dataset = []
    print("start creating dataset")
    for data in dataset_qrel:

        if data["query-id"] in dataset_queries and data["corpus-id"] in dataset_corpus:
            dataset.append({
                "question": dataset_queries[data["query-id"]],
                "passage_positive": [dataset_corpus[data["corpus-id"]]],
                "passage_negative": [],
                "passage_negative_random": [],
            })

    print(f"length of dataset: {len(dataset)}")
    with open(f"./research_notebook/data/mci/{name.split('/')[-1]}_v2.json", "w") as f:
        json.dump(dataset, f, indent=4, ensure_ascii=False)