# -*- 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, 'sys.path.append', 'sys.path.append', False, 'import sys\n'), (20, 'random.seed', 'random.seed', False, 'import random\n'), (70, 'random.sample', 'random.sample', False, 'import random\n'), (76, 'similarities.BM25Similarity', 'BM25Similarity', False, 'from similarities import BM25Similarity\n'), (100, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (102, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (111, 'similarities.evaluation.evaluate', 'evaluate', False, 'from similarities.evaluation import evaluate\n'), (112, 'loguru.logger.info', 'logger.info', False, 'from loguru import logger\n'), (22, 'os.path.realpath', 'os.path.realpath', False, 'import os\n'), (29, 'os.path.join', 'os.path.join', False, 'import os\n'), (35, 'os.path.join', 'os.path.join', False, 'import os\n'), (42, 'os.path.join', 'os.path.join', False, 'import os\n'), (48, 'os.path.join', 'os.path.join', False, 'import os\n'), (83, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (85, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (30, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (31, 'loguru.logger.info', 'logger.info', False, 'from loguru import logger\n'), (32, 'similarities.utils.http_get', 'http_get', False, 'from similarities.utils import http_get\n'), (34, 'loguru.logger.info', 'logger.info', False, 'from loguru import logger\n'), (43, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (44, 'loguru.logger.info', 'logger.info', False, 'from loguru import logger\n'), (45, 'similarities.utils.http_get', 'http_get', False, 'from similarities.utils import http_get\n'), (47, 'loguru.logger.info', 'logger.info', False, 'from loguru import logger\n'), (54, 'similarities.data_loader.SearchDataLoader', 'SearchDataLoader', False, 'from similarities.data_loader import SearchDataLoader\n')]