#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""
 Command line tool to get dense results and validate them
"""

import argparse
import os
import csv
import glob
import json
import gzip
import logging
import pickle
import time
from tqdm import trange
from typing import List, Tuple, Dict, Iterator

import numpy as np
import torch
from torch import nn

from dpr.data.qa_validation import calculate_matches
from dpr.models import init_biencoder_components
from dpr.options import (add_encoder_params, setup_args_gpu, print_args, set_encoder_params_from_state,
                         add_tokenizer_params, add_cuda_params)
from dpr.utils.data_utils import Tensorizer
from dpr.utils.model_utils import setup_for_distributed_mode, get_model_obj, load_states_from_checkpoint
from dpr.indexer.faiss_indexers import DenseIndexer, DenseHNSWFlatIndexer, DenseFlatIndexer

logger = logging.getLogger()
logger.setLevel(logging.INFO)
if logger.hasHandlers():
    logger.handlers.clear()
console = logging.StreamHandler()
logger.addHandler(console)


class DenseRetriever(object):
    """Does passage retrieving over the provided index and question encoder
    """

    def __init__(self, query_encoder: nn.Module, batch_size: int, tensorizer: Tensorizer, index: DenseIndexer):
        self.query_encoder = query_encoder
        self.batch_size = batch_size
        self.tensorizer = tensorizer
        self.index = index

    def generate_query_vectors(self, queries: List[str]) -> torch.Tensor:
        n = len(queries)
        bsz = self.batch_size
        query_vectors = []

        self.query_encoder.eval()
        with torch.no_grad():
            for batch_start in trange(0, n, bsz):
                batch_token_tensors = [self.tensorizer.text_to_tensor(q)
                                       for q in queries[batch_start:batch_start + bsz]]

                q_ids_batch = torch.stack(batch_token_tensors, dim=0).cuda()
                q_seg_batch = torch.zeros_like(q_ids_batch).cuda()
                q_attn_mask = self.tensorizer.get_attn_mask(q_ids_batch)
                _, out, _ = self.query_encoder(q_ids_batch, q_seg_batch, q_attn_mask)

                query_vectors.append(out.cpu())

        query_tensor = torch.cat(query_vectors, dim=0)
        logger.info(f'Total encoded queries tensor {query_tensor.size()}')
        assert query_tensor.size(0) == n

        return query_tensor

    def get_top_docs(self, query_vectors: np.array, top_docs: int = 100) -> List[Tuple[List[object], List[float]]]:
        """Does the retrieval of the best matching passages given the query vectors batch
        """
        time0 = time.time()
        results = self.index.search_knn(query_vectors, top_docs)
        logger.info('index search time: %f sec.', time.time() - time0)
        return results


def parse_qa_csv_file(location) -> Iterator[Tuple[str, List[str]]]:
    with open(location) as fin:
        reader = csv.reader(fin, delimiter='\t')
        for row in reader:
            question = row[0]
            answers = eval(row[1])
            yield question, answers


def load_passages(ctx_file: str) -> Dict[object, Tuple[str, str]]:
    docs = {}
    logger.info(f'Reading passages from {ctx_file}')
    if ctx_file.endswith(".gz"):
        with gzip.open(ctx_file, 'rt') as tsv_file:
            # file format: doc_id, doc_text, title
            num_field = None
            for line in tsv_file:
                segs = line.strip().split('\t')
                pid, text, title = segs[:3]
                if pid != 'id':
                    if pid in docs:
                        logger.warning(f'Duplicate paragraph: {pid}')
                    docs[pid] = (text, title)
                else:
                    num_field = len(segs)
                if len(segs) != num_field:
                    logger.warning(f'Wrong line format: {pid}')
    else:
        with open(ctx_file) as tsv_file:
            # file format: doc_id, doc_text, title
            num_field = None
            for line in tsv_file:
                segs = line.strip().split('\t')
                pid, text, title = segs[:3]
                if pid != 'id':
                    if pid in docs:
                        logger.warning(f'Duplicate paragraph: {pid}')
                    docs[pid] = (text, title)
                else:
                    num_field = len(segs)
                if len(segs) != num_field:
                    logger.warning(f'Wrong line format: {pid}')
    return docs


def validate(passages: Dict[object, Tuple[str, str]],
             answers: List[List[str]],
             result_ctx_ids: List[Tuple[List[object], List[float]]],
             workers_num: int, match_type: str) -> List[List[bool]]:
    match_stats = calculate_matches(passages, answers, result_ctx_ids, workers_num, match_type)
    top_k_hits = match_stats.top_k_hits

    logger.info('Validation results: top k documents hits %s', top_k_hits)
    top_k_hits = [v / len(result_ctx_ids) for v in top_k_hits]
    logger.info('Validation results: top k documents hits accuracy %s', top_k_hits)
    return match_stats.questions_doc_hits


def save_results(passages: Dict[object, Tuple[str, str]],
                 questions: List[str],
                 answers: List[List[str]],
                 top_passages_and_scores: List[Tuple[List[object], List[float]]],
                 per_question_hits: List[List[bool]],
                 out_file: str):
    # join passages text with the result ids, their questions and assigning has|no answer labels
    merged_data = []
    assert len(per_question_hits) == len(questions) == len(answers)
    for i, q in enumerate(questions):
        q_answers = answers[i]
        results_and_scores = top_passages_and_scores[i]
        hits = per_question_hits[i]
        docs = [passages[doc_id] for doc_id in results_and_scores[0]]
        scores = [str(score) for score in results_and_scores[1]]
        ctxs_num = len(hits)

        merged_data.append({
            'question': q,
            'answers': q_answers,
            'ctxs': [
                {
                    'id': results_and_scores[0][c],
                    'title': docs[c][1],
                    'text': docs[c][0],
                    'score': scores[c],
                    'has_answer': hits[c],
                } for c in range(ctxs_num)
            ]
        })

    with open(out_file, "w") as writer:
        writer.write(json.dumps(merged_data, indent=2) + "\n")
    logger.info('Saved results * scores  to %s', out_file)


def iterate_encoded_files(vector_files: list) -> Iterator[Tuple[object, np.array]]:
    for i, file in enumerate(vector_files):
        logger.info(f'Reading file {file}')
        with open(file, "rb") as reader:
            doc_vectors = pickle.load(reader)
            for db_id, doc_vector in doc_vectors:
                yield db_id, doc_vector


def main():
    saved_state = load_states_from_checkpoint(args.model_file)
    set_encoder_params_from_state(saved_state.encoder_params, args)

    tensorizer, encoder, _ = init_biencoder_components(args.encoder_model_type, args, inference_only=True)

    encoder = encoder.question_model

    encoder, _ = setup_for_distributed_mode(encoder, None, args.device, args.n_gpu,
                                            args.local_rank, args.fp16)
    encoder.eval()

    # load weights from the model file
    model_to_load = get_model_obj(encoder)
    logger.info('Loading saved model state ...')

    prefix_len = len('question_model.')
    question_encoder_state = {key[prefix_len:]: value
                              for (key, value) in saved_state.model_dict.items() if key.startswith('question_model.')}
    model_to_load.load_state_dict(question_encoder_state)
    vector_size = model_to_load.get_out_size()
    logger.info(f'Encoder vector_size={vector_size}')

    if args.hnsw_index:
        index = DenseHNSWFlatIndexer(vector_size, args.index_buffer)
    else:
        index = DenseFlatIndexer(vector_size, args.index_buffer)

    retriever = DenseRetriever(encoder, args.batch_size, tensorizer, index)

    # index all passages
    ctx_files_pattern = args.encoded_ctx_file
    dense_files = sorted(glob.glob(ctx_files_pattern))
    index_prefix = "_".join(dense_files[0].split("_")[:-1]) if not args.save_or_load_index else args.save_or_load_index
    if args.save_or_load_index and (os.path.exists(index_prefix) or os.path.exists(index_prefix + ".index.dpr")):
        retriever.index.deserialize_from(index_prefix)
    else:
        logger.info(f'Reading all passage vectors from: {", ".join(dense_files)}')
        retriever.index.index_data(dense_files)
        if args.save_or_load_index:
            retriever.index.serialize(index_prefix)

    # get questions & answers
    question_list = []
    answers_list = []
    for ds_item in parse_qa_csv_file(args.qa_file):
        question, answers = ds_item
        question_list.append(question)
        answers_list.append(answers)

    # encode questions to vectors
    questions_tensor = retriever.generate_query_vectors(question_list)

    # get top k results
    top_ids_and_scores = retriever.get_top_docs(questions_tensor.numpy(), args.n_docs)

    # validate
    all_passages = load_passages(args.ctx_file)
    if len(all_passages) == 0:
        raise RuntimeError('No passages data found. Please specify ctx_file param properly.')
    questions_doc_hits = validate(all_passages, answers_list, top_ids_and_scores,
                                  args.validation_workers, args.match)
    if args.out_file:
        save_results(all_passages, question_list, answers_list, top_ids_and_scores, questions_doc_hits, args.out_file)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    add_encoder_params(parser)
    add_tokenizer_params(parser)
    add_cuda_params(parser)

    parser.add_argument('--qa_file', required=True, type=str, default=None,
                        help="Question and answers file of the format: question \\t ['answer1','answer2', ...]")
    parser.add_argument('--ctx_file', required=True, type=str, default=None,
                        help="All passages file in the tsv format: id \\t passage_text \\t title")
    parser.add_argument('--encoded_ctx_file', type=str, default=None,
                        help='Glob path to encoded passages (from generate_dense_embeddings tool)')

    parser.add_argument('--batch_size', type=int, default=32,
                        help="Batch size for question encoder forward pass")

    parser.add_argument('--n-docs', type=int, default=200, help="Amount of top docs to return")
    parser.add_argument('--index_buffer', type=int, default=50000,
                        help="Temporal memory data buffer size (in samples) for indexer")
    parser.add_argument("--hnsw_index", action='store_true',
                        help='If enabled, use inference time efficient HNSW index')
    parser.add_argument("--save_or_load_index", type=str, default=None,
                        help='Index file path (prefix) to load/save')

    parser.add_argument('--validation_workers', type=int, default=16,
                        help="Number of parallel processes to validate results")
    parser.add_argument('--match', type=str, default='string', choices=['regex', 'string'],
                        help="Answer matching logic type")
    parser.add_argument('--out_file', type=str, default=None,
                        help='Output .json file path to write results to')

    args = parser.parse_args()

    assert args.model_file, 'Please specify --model_file checkpoint to init model weights'

    setup_args_gpu(args)
    print_args(args)
    main()
