# -*- 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}")

[((14, 0, 14, 24), 'sys.path.append', 'sys.path.append', ({(14, 16, 14, 23): '"""../.."""'}, {}), "('../..')", False, 'import sys\n'), ((20, 0, 20, 15), 'random.seed', 'random.seed', ({(20, 12, 20, 14): '(42)'}, {}), '(42)', False, 'import random\n'), ((70, 17, 70, 56), 'random.sample', 'random.sample', ({(70, 31, 70, 47): 'remaining_corpus', (70, 49, 70, 55): 'sample'}, {}), '(remaining_corpus, sample)', False, 'import random\n'), ((76, 8, 76, 35), 'similarities.BM25Similarity', 'BM25Similarity', ({(76, 23, 76, 34): 'corpus_docs'}, {}), '(corpus_docs)', False, 'from similarities import BM25Similarity\n'), ((100, 8, 100, 31), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((102, 6, 102, 29), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((111, 32, 111, 56), 'similarities.evaluation.evaluate', 'evaluate', ({(111, 41, 111, 46): 'qrels', (111, 48, 111, 55): 'results'}, {}), '(qrels, results)', False, 'from similarities.evaluation import evaluate\n'), ((112, 0, 112, 27), 'loguru.logger.info', 'logger.info', ({(112, 12, 112, 26): 'f"""MAP: {_map}"""'}, {}), "(f'MAP: {_map}')", False, 'from loguru import logger\n'), ((22, 27, 22, 53), 'os.path.realpath', 'os.path.realpath', ({(22, 44, 22, 52): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((29, 15, 29, 52), 'os.path.join', 'os.path.join', ({(29, 28, 29, 36): 'pwd_path', (29, 38, 29, 51): '"""scifact.zip"""'}, {}), "(pwd_path, 'scifact.zip')", False, 'import os\n'), ((35, 16, 35, 47), 'os.path.join', 'os.path.join', ({(35, 29, 35, 37): 'pwd_path', (35, 39, 35, 46): 'dataset'}, {}), '(pwd_path, dataset)', False, 'import os\n'), ((42, 15, 42, 59), 'os.path.join', 'os.path.join', ({(42, 28, 42, 36): 'pwd_path', (42, 38, 42, 58): '"""dbpedia-entity.zip"""'}, {}), "(pwd_path, 'dbpedia-entity.zip')", False, 'import os\n'), ((48, 16, 48, 47), 'os.path.join', 'os.path.join', ({(48, 29, 48, 37): 'pwd_path', (48, 39, 48, 46): 'dataset'}, {}), '(pwd_path, dataset)', False, 'import os\n'), ((83, 12, 83, 35), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((85, 10, 85, 33), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((30, 11, 30, 35), 'os.path.exists', 'os.path.exists', ({(30, 26, 30, 34): 'zip_file'}, {}), '(zip_file)', False, 'import os\n'), ((31, 8, 31, 57), 'loguru.logger.info', 'logger.info', ({(31, 20, 31, 56): '"""Dataset not exists, downloading..."""'}, {}), "('Dataset not exists, downloading...')", False, 'from loguru import logger\n'), ((32, 8, 32, 45), 'similarities.utils.http_get', 'http_get', (), '', False, 'from similarities.utils import http_get\n'), ((34, 8, 34, 65), 'loguru.logger.info', 'logger.info', ({(34, 20, 34, 64): '"""Dataset already exists, skipping download."""'}, {}), "('Dataset already exists, skipping download.')", False, 'from loguru import logger\n'), ((43, 11, 43, 35), 'os.path.exists', 'os.path.exists', ({(43, 26, 43, 34): 'zip_file'}, {}), '(zip_file)', False, 'import os\n'), ((44, 8, 44, 57), 'loguru.logger.info', 'logger.info', ({(44, 20, 44, 56): '"""Dataset not exists, downloading..."""'}, {}), "('Dataset not exists, downloading...')", False, 'from loguru import logger\n'), ((45, 8, 45, 45), 'similarities.utils.http_get', 'http_get', (), '', False, 'from similarities.utils import http_get\n'), ((47, 8, 47, 65), 'loguru.logger.info', 'logger.info', ({(47, 20, 47, 64): '"""Dataset already exists, skipping download."""'}, {}), "('Dataset already exists, skipping download.')", False, 'from loguru import logger\n'), ((54, 25, 54, 52), 'similarities.data_loader.SearchDataLoader', 'SearchDataLoader', ({(54, 42, 54, 51): 'data_path'}, {}), '(data_path)', False, 'from similarities.data_loader import SearchDataLoader\n')]