# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description:
"""
import datetime
import os
import pathlib
import random
import sys

from loguru import logger

sys.path.append('../..')
from similarities import BM25Similarity
from similarities.utils import http_get
from similarities.data_loader import SearchDataLoader
from similarities.evaluation import evaluate

random.seed(42)

pwd_path = os.path.dirname(os.path.realpath(__file__))


def get_scifact():
    # Download scifact.zip dataset and unzip the dataset
    dataset = "scifact"
    url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip".format(dataset)
    zip_file = os.path.join(pwd_path, "scifact.zip")
    if not os.path.exists(zip_file):
        logger.info("Dataset not exists, downloading...")
        http_get(url, zip_file, extract=True)
    else:
        logger.info("Dataset already exists, skipping download.")
    data_path = os.path.join(pwd_path, dataset)
    return data_path


def get_dbpedia():
    dataset = "dbpedia-entity"
    url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip".format(dataset)
    zip_file = os.path.join(pwd_path, "dbpedia-entity.zip")
    if not os.path.exists(zip_file):
        logger.info("Dataset not exists, downloading...")
        http_get(url, zip_file, extract=True)
    else:
        logger.info("Dataset already exists, skipping download.")
    data_path = os.path.join(pwd_path, dataset)
    return data_path


data_path = get_scifact()
#### Loading test queries and corpus in DBPedia
corpus, queries, qrels = SearchDataLoader(data_path).load(split="test")
corpus_ids, query_ids = list(corpus), list(queries)
logger.info(f"corpus: {len(corpus)}, queries: {len(queries)}")

#### Randomly sample 1M pairs from Original Corpus (4.63M pairs) 
#### First include all relevant documents (i.e. present in qrels)
corpus_set = set()
for query_id in qrels:
    corpus_set.update(list(qrels[query_id].keys()))
corpus_new = {corpus_id: corpus[corpus_id] for corpus_id in corpus_set}

#### Remove already seen k relevant documents and sample (1M - k) docs randomly
remaining_corpus = list(set(corpus_ids) - corpus_set)
sample = min(1000000 - len(corpus_set), len(remaining_corpus))
# sample = 10

for corpus_id in random.sample(remaining_corpus, sample):
    corpus_new[corpus_id] = corpus[corpus_id]

corpus_docs = {corpus_id: corpus_new[corpus_id]['title'] + corpus_new[corpus_id]['text'] for corpus_id, corpus in
               corpus_new.items()}
#### Index 1M passages into the index (seperately)
model = BM25Similarity(corpus_docs)
#### Saving benchmark times
time_taken_all = {}
for query_id in query_ids:
    query = {query_id: queries[query_id]}

    #### Measure time to retrieve top-10 BM25 documents using single query latency
    start = datetime.datetime.now()
    q_res = model.most_similar(query, topn=10)
    end = datetime.datetime.now()
    # print(q_res)

    #### Measuring time taken in ms (milliseconds)
    time_taken = (end - start)
    time_taken = time_taken.total_seconds() * 1000
    time_taken_all[query_id] = time_taken
    # logger.info("query: {}: {} {:.2f}ms".format(query_id, query, time_taken))
    # logger.info("\tsearch result: {}".format(results[:2]))

time_taken = list(time_taken_all.values())
logger.info("Average time taken: {:.2f}ms".format(sum(time_taken) / len(time_taken_all)))

#### Saving benchmark times with batch
# queries = [queries[query_id] for query_id in query_ids]
start = datetime.datetime.now()
results = model.most_similar(queries, topn=10)
end = datetime.datetime.now()
#### Measuring time taken in ms (milliseconds)
time_taken = (end - start)
time_taken = time_taken.total_seconds() * 1000
logger.info("All, Spend {:.2f}ms".format(time_taken))
logger.info("Average time taken: {:.2f}ms".format(time_taken / len(queries)))
logger.info(f"Results size: {len(results)}")

#### Evaluate your retrieval using NDCG@k, MAP@K ...
ndcg, _map, recall, precision = evaluate(qrels, results)
logger.info(f"MAP: {_map}")

[((173, 197), 'sys.path.append', 'sys.path.append', (['"""../.."""'], {}), "('../..')\n", (188, 197), False, 'import sys\n'), ((378, 393), 'random.seed', 'random.seed', (['(42)'], {}), '(42)\n', (389, 393), False, 'import random\n'), ((2285, 2324), 'random.sample', 'random.sample', (['remaining_corpus', 'sample'], {}), '(remaining_corpus, sample)\n', (2298, 2324), False, 'import random\n'), ((2581, 2608), 'similarities.BM25Similarity', 'BM25Similarity', (['corpus_docs'], {}), '(corpus_docs)\n', (2595, 2608), False, 'from similarities import BM25Similarity\n'), ((3503, 3526), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (3524, 3526), False, 'import datetime\n'), ((3580, 3603), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (3601, 3603), False, 'import datetime\n'), ((3988, 4012), 'similarities.evaluation.evaluate', 'evaluate', (['qrels', 'results'], {}), '(qrels, results)\n', (3996, 4012), False, 'from similarities.evaluation import evaluate\n'), ((4013, 4040), 'loguru.logger.info', 'logger.info', (['f"""MAP: {_map}"""'], {}), "(f'MAP: {_map}')\n", (4024, 4040), False, 'from loguru import logger\n'), ((422, 448), 'os.path.realpath', 'os.path.realpath', (['__file__'], {}), '(__file__)\n', (438, 448), False, 'import os\n'), ((669, 706), 'os.path.join', 'os.path.join', (['pwd_path', '"""scifact.zip"""'], {}), "(pwd_path, 'scifact.zip')\n", (681, 706), False, 'import os\n'), ((940, 971), 'os.path.join', 'os.path.join', (['pwd_path', 'dataset'], {}), '(pwd_path, dataset)\n', (952, 971), False, 'import os\n'), ((1162, 1206), 'os.path.join', 'os.path.join', (['pwd_path', '"""dbpedia-entity.zip"""'], {}), "(pwd_path, 'dbpedia-entity.zip')\n", (1174, 1206), False, 'import os\n'), ((1440, 1471), 'os.path.join', 'os.path.join', (['pwd_path', 'dataset'], {}), '(pwd_path, dataset)\n', (1452, 1471), False, 'import os\n'), ((2822, 2845), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (2843, 2845), False, 'import datetime\n'), ((2903, 2926), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (2924, 2926), False, 'import datetime\n'), ((718, 742), 'os.path.exists', 'os.path.exists', (['zip_file'], {}), '(zip_file)\n', (732, 742), False, 'import os\n'), ((752, 801), 'loguru.logger.info', 'logger.info', (['"""Dataset not exists, downloading..."""'], {}), "('Dataset not exists, downloading...')\n", (763, 801), False, 'from loguru import logger\n'), ((810, 847), 'similarities.utils.http_get', 'http_get', (['url', 'zip_file'], {'extract': '(True)'}), '(url, zip_file, extract=True)\n', (818, 847), False, 'from similarities.utils import http_get\n'), ((866, 923), 'loguru.logger.info', 'logger.info', (['"""Dataset already exists, skipping download."""'], {}), "('Dataset already exists, skipping download.')\n", (877, 923), False, 'from loguru import logger\n'), ((1218, 1242), 'os.path.exists', 'os.path.exists', (['zip_file'], {}), '(zip_file)\n', (1232, 1242), False, 'import os\n'), ((1252, 1301), 'loguru.logger.info', 'logger.info', (['"""Dataset not exists, downloading..."""'], {}), "('Dataset not exists, downloading...')\n", (1263, 1301), False, 'from loguru import logger\n'), ((1310, 1347), 'similarities.utils.http_get', 'http_get', (['url', 'zip_file'], {'extract': '(True)'}), '(url, zip_file, extract=True)\n', (1318, 1347), False, 'from similarities.utils import http_get\n'), ((1366, 1423), 'loguru.logger.info', 'logger.info', (['"""Dataset already exists, skipping download."""'], {}), "('Dataset already exists, skipping download.')\n", (1377, 1423), False, 'from loguru import logger\n'), ((1594, 1621), 'similarities.data_loader.SearchDataLoader', 'SearchDataLoader', (['data_path'], {}), '(data_path)\n', (1610, 1621), False, 'from similarities.data_loader import SearchDataLoader\n')]