"""
- Script to generate predictions for long answer type questions in simplified validation dataset
  (v1.0-simplified_nq-dev-all.jsonl) and save the predictions to a predictions.json file
- Pass arguments using command line
- Restricted for device with CUDA
- Make sure the loaded config, tokenizer and weights belong to the same model
- Runs with the model trained with the train_v0.py script
"""
import os
import json
import logging
import argparse
import jsonlines

import numpy as np
from tqdm import tqdm
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertConfig

from custom_typing.datasets import SimplifiedNaturalQADataset
from models.bert.bert_for_qa import BertForQuestionAnswering
from utils.collator import CollatorForValidation

logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S')


def parse_data_from_json_file(val_dataset: str, max_data: int = 1e10):
    """Reads and parses the json file for simplified natural questions validation dataset

    Parsing training examples and creating a data dictionary of examples (positive and negative).
    Negative examples are generated by uniform sampling from the list of incorrect long answer candidates for positive
    examples

    Args:
        val_dataset (str): path to simplified training dataset jsonl file
        max_data (int): max number of examples to validate, majorly used for debugging purposes

    Returns:
        id_list (List[int]): list of document ids
        id_candidate_list_sorted (List(Tuple(int, int))):sorted list of (document id, long answer candidate index) based
            on length of long answer candidate
        data_dict (dict):
    """
    # check if input file is of type jsonl
    assert os.path.splitext(val_dataset)[-1] == ".jsonl", "dataset file type is not jsonl, check the file provided"
    # store all document ids
    id_list = []
    # store all candidates
    id_candidate_list = []
    # store length of all candidates
    id_candidate_len_list = []
    # store data id to candidate length dict
    id_candidate_len_dict = {}
    data_dict = {}

    with jsonlines.open(val_dataset) as reader:
        for n, data_line in enumerate(tqdm(reader)):
            if n > max_data:
                break
            doc_id = data_line['example_id']
            id_list.append(doc_id)

            # initialize data_dict
            data_dict[doc_id] = {
                'document_text': ' '.join([item['token'] for item in data_line['document_tokens']]),
                'question_text': data_line['question_text'],
                'long_answer_candidates': data_line['long_answer_candidates'],
            }

            question_len = len(data_line['question_text'].split())

            # We use the white space tokenized version to estimate candidate length here.
            for i, candidate in enumerate(data_line['long_answer_candidates']):
                id_candidate_list.append((doc_id, i))
                candidate_length = question_len + candidate['end_token'] - candidate['start_token']
                id_candidate_len_list.append(candidate_length)
                id_candidate_len_dict[(doc_id, i)] = candidate_length

    # sorting candidates based on candidate's length
    sorted_index = np.argsort(np.array(id_candidate_len_list))
    id_candidate_list_sorted = []
    for i in range(len(id_candidate_list)):
        id_candidate_list_sorted.append(id_candidate_list[sorted_index[i]])

    return id_list, id_candidate_list_sorted, data_dict


if __name__=='__main__':
    parser = argparse.ArgumentParser(description="parser to mine hard examples from a set of examples")
    parser.add_argument('-d', '--val_dataset', help='path to dataset examples json file', type=str,
                        default='../datasets/natural_questions_simplified/v1.0-simplified_nq-dev-all.jsonl')
    parser.add_argument('-o', '--output_path', help='path to store predictions', type=str,
                        default='../predictions/bert_base_uncased/')
    parser.add_argument('-m', '--model_path', help='path to a saved model', type=str, default='bert-base-uncased')
    parser.add_argument('-w', '--weights', help='path to saved weights for the model', type=str,
                        default='../weights/bert-base-uncased/epoch1/')
    parser.add_argument('--fp16', action='store_true', help='mention if loaded model is trained on half precision')
    args = parser.parse_args()

    logging.info("parsing validation dataset")
    id_list, id_candidate_list_sorted, data_dict = parse_data_from_json_file(args.val_dataset)

    # hyperparameters
    max_seq_length = 384
    max_question_length = 64
    batch_size = 768

    # load model
    logging.info("loading config")
    config = BertConfig.from_pretrained(args.model_path)
    config.num_labels = 5
    tokenizer = BertTokenizer.from_pretrained(args.model_path, do_lower_case=True)
    model = BertForQuestionAnswering.from_pretrained(args.weights, config=config)

    if torch.cuda.is_available():
        model.cuda()
    logging.info(f"fp16: {args.fp16}")
    if args.fp16:
        from apex import amp
        model = amp.initialize(model, opt_level="O1", verbosity=0)
    if torch.cuda.device_count() > 1:
        model = torch.nn.DataParallel(model)

    # testing
    # iterator for testing
    test_data_generator = SimplifiedNaturalQADataset(id_list=id_candidate_list_sorted)
    test_collator = CollatorForValidation(data_dict=data_dict,
                            tokenizer=tokenizer,
                            max_seq_length=max_seq_length,
                            max_question_length=max_question_length)
    test_generator = DataLoader(dataset=test_data_generator,
                                collate_fn=test_collator,
                                batch_size=batch_size,
                                shuffle=False,
                                num_workers=16,
                                pin_memory=True)

    # evaluating model on dataset
    logging.info("Evaluating")
    model.eval()

    classifier_probs = np.zeros((len(id_candidate_list_sorted), 5), dtype=np.float32)  # class

    for j, (batch_input_ids, batch_attention_mask, batch_token_type_ids) in enumerate(tqdm(test_generator)):
        with torch.no_grad():
            start = j*batch_size
            end = start + batch_size
            if j == len(test_generator) - 1:
                end = len(test_generator.dataset)
            if torch.cuda.is_available():
                batch_input_ids = batch_input_ids.cuda()
                batch_attention_mask = batch_attention_mask.cuda()
                batch_token_type_ids = batch_token_type_ids.cuda()
            start_position_logits, end_position_logits, classifier_logits = model(batch_input_ids, batch_attention_mask, batch_token_type_ids)
            classifier_probs[start:end] += nn.functional.softmax(classifier_logits, dim=1).cpu().data.numpy()

    # Processing long answers:
    # initialize a temporary dictionary to store prediction values.
    temp_dict = {}
    for doc_id in id_list:
        temp_dict[doc_id] = {
            'long_answer': {'start_token': -1, 'end_token': -1},
            'long_answer_score': -1.0,
            'short_answers': [{'start_token': -1, 'end_token': -1}],
            'short_answers_score': -1.0,
            'yes_no_answer': 'NONE'
        }

    # from candidates to document
    for i, (doc_id, candidate_index) in enumerate(tqdm(id_candidate_list_sorted)):
        # process long answer
        la_candidate = data_dict[doc_id]['long_answer_candidates'][candidate_index]
        long_answer_score = 1.0 - classifier_probs[i,0] # 1- no_answer_score
        if long_answer_score > temp_dict[doc_id]['long_answer_score']:
            temp_dict[doc_id]['long_answer_score'] = long_answer_score
            temp_dict[doc_id]['long_answer']['start_token'] = la_candidate['start_token']
            temp_dict[doc_id]['long_answer']['end_token'] = la_candidate['end_token']

    # Copy the temporary dictionary into the final dictionary that meets the required format for validation.
    final_dict = {'predictions': []}
    for doc_id in id_list:
        prediction_dict = {
            'example_id': doc_id,
            'long_answer': {'start_byte': -1, 'end_byte': -1,
                            'start_token': temp_dict[doc_id]['long_answer']['start_token'],
                            'end_token': temp_dict[doc_id]['long_answer']['end_token']},
            'long_answer_score': temp_dict[doc_id]['long_answer_score'],
            'short_answers': [{'start_byte': -1, 'end_byte': -1, 'start_token': -1, 'end_token': -1}],
            'short_answers_score': -1.0,
            'yes_no_answer': 'NONE'
        }
        final_dict['predictions'].append(prediction_dict)

    # make sure output directory exists
    if not os.path.exists(args.output_path):
        os.makedirs(args.output_path, exist_ok=True)
    logging.info(f"writing predictions to {args.output_path}")
    # write to json file
    with open(os.path.join(args.output_path, 'predictions.json'), 'w') as f:
        json.dump(final_dict, f)

"""
- Script to generate predictions for long answer type questions in simplified validation dataset
  (v1.0-simplified_nq-dev-all.jsonl) and save the predictions to a predictions.json file
- Pass arguments using command line
- Restricted for device with CUDA
- Make sure the loaded config, tokenizer and weights belong to the same model
- Runs with the model trained with the train_v0.py script
"""
import os
import json
import logging
import argparse
import jsonlines

import numpy as np
from tqdm import tqdm
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertConfig

from custom_typing.datasets import SimplifiedNaturalQADataset
from models.bert.bert_for_qa import BertForQuestionAnswering
from utils.collator import CollatorForValidation

logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S')


def parse_data_from_json_file(val_dataset: str, max_data: int = 1e10):
    """Reads and parses the json file for simplified natural questions validation dataset

    Parsing training examples and creating a data dictionary of examples (positive and negative).
    Negative examples are generated by uniform sampling from the list of incorrect long answer candidates for positive
    examples

    Args:
        val_dataset (str): path to simplified training dataset jsonl file
        max_data (int): max number of examples to validate, majorly used for debugging purposes

    Returns:
        id_list (List[int]): list of document ids
        id_candidate_list_sorted (List(Tuple(int, int))):sorted list of (document id, long answer candidate index) based
            on length of long answer candidate
        data_dict (dict):
    """
    # check if input file is of type jsonl
    assert os.path.splitext(val_dataset)[-1] == ".jsonl", "dataset file type is not jsonl, check the file provided"
    # store all document ids
    id_list = []
    # store all candidates
    id_candidate_list = []
    # store length of all candidates
    id_candidate_len_list = []
    # store data id to candidate length dict
    id_candidate_len_dict = {}
    data_dict = {}

    with jsonlines.open(val_dataset) as reader:
        for n, data_line in enumerate(tqdm(reader)):
            if n > max_data:
                break
            doc_id = data_line['example_id']
            id_list.append(doc_id)

            # initialize data_dict
            data_dict[doc_id] = {
                'document_text': ' '.join([item['token'] for item in data_line['document_tokens']]),
                'question_text': data_line['question_text'],
                'long_answer_candidates': data_line['long_answer_candidates'],
            }

            question_len = len(data_line['question_text'].split())

            # We use the white space tokenized version to estimate candidate length here.
            for i, candidate in enumerate(data_line['long_answer_candidates']):
                id_candidate_list.append((doc_id, i))
                candidate_length = question_len + candidate['end_token'] - candidate['start_token']
                id_candidate_len_list.append(candidate_length)
                id_candidate_len_dict[(doc_id, i)] = candidate_length

    # sorting candidates based on candidate's length
    sorted_index = np.argsort(np.array(id_candidate_len_list))
    id_candidate_list_sorted = []
    for i in range(len(id_candidate_list)):
        id_candidate_list_sorted.append(id_candidate_list[sorted_index[i]])

    return id_list, id_candidate_list_sorted, data_dict


if __name__=='__main__':
    parser = argparse.ArgumentParser(description="parser to mine hard examples from a set of examples")
    parser.add_argument('-d', '--val_dataset', help='path to dataset examples json file', type=str,
                        default='../datasets/natural_questions_simplified/v1.0-simplified_nq-dev-all.jsonl')
    parser.add_argument('-o', '--output_path', help='path to store predictions', type=str,
                        default='../predictions/bert_base_uncased/')
    parser.add_argument('-m', '--model_path', help='path to a saved model', type=str, default='bert-base-uncased')
    parser.add_argument('-w', '--weights', help='path to saved weights for the model', type=str,
                        default='../weights/bert-base-uncased/epoch1/')
    parser.add_argument('--fp16', action='store_true', help='mention if loaded model is trained on half precision')
    args = parser.parse_args()

    logging.info("parsing validation dataset")
    id_list, id_candidate_list_sorted, data_dict = parse_data_from_json_file(args.val_dataset)

    # hyperparameters
    max_seq_length = 384
    max_question_length = 64
    batch_size = 768

    # load model
    logging.info("loading config")
    config = BertConfig.from_pretrained(args.model_path)
    config.num_labels = 5
    tokenizer = BertTokenizer.from_pretrained(args.model_path, do_lower_case=True)
    model = BertForQuestionAnswering.from_pretrained(args.weights, config=config)

    if torch.cuda.is_available():
        model.cuda()
    logging.info(f"fp16: {args.fp16}")
    if args.fp16:
        from apex import amp
        model = amp.initialize(model, opt_level="O1", verbosity=0)
    if torch.cuda.device_count() > 1:
        model = torch.nn.DataParallel(model)

    # testing
    # iterator for testing
    test_data_generator = SimplifiedNaturalQADataset(id_list=id_candidate_list_sorted)
    test_collator = CollatorForValidation(data_dict=data_dict,
                            tokenizer=tokenizer,
                            max_seq_length=max_seq_length,
                            max_question_length=max_question_length)
    test_generator = DataLoader(dataset=test_data_generator,
                                collate_fn=test_collator,
                                batch_size=batch_size,
                                shuffle=False,
                                num_workers=16,
                                pin_memory=True)

    # evaluating model on dataset
    logging.info("Evaluating")
    model.eval()

    classifier_probs = np.zeros((len(id_candidate_list_sorted), 5), dtype=np.float32)  # class

    for j, (batch_input_ids, batch_attention_mask, batch_token_type_ids) in enumerate(tqdm(test_generator)):
        with torch.no_grad():
            start = j*batch_size
            end = start + batch_size
            if j == len(test_generator) - 1:
                end = len(test_generator.dataset)
            if torch.cuda.is_available():
                batch_input_ids = batch_input_ids.cuda()
                batch_attention_mask = batch_attention_mask.cuda()
                batch_token_type_ids = batch_token_type_ids.cuda()
            start_position_logits, end_position_logits, classifier_logits = model(batch_input_ids, batch_attention_mask, batch_token_type_ids)
            classifier_probs[start:end] += nn.functional.softmax(classifier_logits, dim=1).cpu().data.numpy()

    # Processing long answers:
    # initialize a temporary dictionary to store prediction values.
    temp_dict = {}
    for doc_id in id_list:
        temp_dict[doc_id] = {
            'long_answer': {'start_token': -1, 'end_token': -1},
            'long_answer_score': -1.0,
            'short_answers': [{'start_token': -1, 'end_token': -1}],
            'short_answers_score': -1.0,
            'yes_no_answer': 'NONE'
        }

    # from candidates to document
    for i, (doc_id, candidate_index) in enumerate(tqdm(id_candidate_list_sorted)):
        # process long answer
        la_candidate = data_dict[doc_id]['long_answer_candidates'][candidate_index]
        long_answer_score = 1.0 - classifier_probs[i,0] # 1- no_answer_score
        if long_answer_score > temp_dict[doc_id]['long_answer_score']:
            temp_dict[doc_id]['long_answer_score'] = long_answer_score
            temp_dict[doc_id]['long_answer']['start_token'] = la_candidate['start_token']
            temp_dict[doc_id]['long_answer']['end_token'] = la_candidate['end_token']

    # Copy the temporary dictionary into the final dictionary that meets the required format for validation.
    final_dict = {'predictions': []}
    for doc_id in id_list:
        prediction_dict = {
            'example_id': doc_id,
            'long_answer': {'start_byte': -1, 'end_byte': -1,
                            'start_token': temp_dict[doc_id]['long_answer']['start_token'],
                            'end_token': temp_dict[doc_id]['long_answer']['end_token']},
            'long_answer_score': temp_dict[doc_id]['long_answer_score'],
            'short_answers': [{'start_byte': -1, 'end_byte': -1, 'start_token': -1, 'end_token': -1}],
            'short_answers_score': -1.0,
            'yes_no_answer': 'NONE'
        }
        final_dict['predictions'].append(prediction_dict)

    # make sure output directory exists
    if not os.path.exists(args.output_path):
        os.makedirs(args.output_path, exist_ok=True)
    logging.info(f"writing predictions to {args.output_path}")
    # write to json file
    with open(os.path.join(args.output_path, 'predictions.json'), 'w') as f:
        json.dump(final_dict, f)

- Script to generate predictions for long answer type questions in simplified validation dataset   (v1.0-simplified_nq-dev-all.jsonl) and save the predictions to a predictions.json file - Pass arguments using command line - Restricted for device with CUDA - Make sure the loaded config, tokenizer and weights belong to the same model - Runs with the model trained with the train_v0.py script Reads and parses the json file for simplified natural questions validation dataset      Parsing training examples and creating a data dictionary of examples (positive and negative).     Negative examples are generated by uniform sampling from the list of incorrect long answer candidates for positive     examples      Args:         val_dataset (str): path to simplified training dataset jsonl file         max_data (int): max number of examples to validate, majorly used for debugging purposes      Returns:         id_list (List[int]): list of document ids         id_candidate_list_sorted (List(Tuple(int, int))):sorted list of (document id, long answer candidate index) based             on length of long answer candidate         data_dict (dict): # check if input file is of type jsonl # store all document ids # store all candidates # store length of all candidates # store data id to candidate length dict # initialize data_dict # We use the white space tokenized version to estimate candidate length here. # sorting candidates based on candidate's length # hyperparameters # load model # testing # iterator for testing # evaluating model on dataset # class # Processing long answers: # initialize a temporary dictionary to store prediction values. # from candidates to document # process long answer # 1- no_answer_score # Copy the temporary dictionary into the final dictionary that meets the required format for validation. # make sure output directory exists # write to json file