# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import collections
from multiprocessing.sharedctypes import Value
import os
import torch
import numpy as np
import glob

def format_multichoice(multichoice_options):

    options_text = ["({}) {}".format(chr(ord('A')+i), option) for i, option in zip(range(len(multichoice_options)), multichoice_options)]
    return "Choose one based on the following options: {}".format(" ".join(options_text))

def format_multichoice_question(question, multichoice_options):

    return  "{}\n{}".format(question, format_multichoice(multichoice_options))

def format_answer(answer):
    return " {}".format(answer)

"""GPT ft dataset."""
def preprocess(data_file, inference_only=False, retrieved_neighbours=False, fix_newsqa=False):

    nq_examples = []
    for my_data_file in sorted(glob.glob(data_file)):
        with open(my_data_file, "r", encoding='utf-8') as f:
            nq_examples.extend(json.load(f))
    
    data = []
    for instance in nq_examples:
        question = instance["question"]
        if 'qa_type' in instance and instance['qa_type'] == "multi_choice_qa":
            question = format_multichoice_question(question, instance["multichoice_options"])
        if True:
            if retrieved_neighbours:
                contexts = instance["ctxs"]
                neighbours = ["title: " + ctx["title"] + ", source: " + ctx["text"] for ctx in contexts] 
            else:
                if "document" in instance:
                    doc = instance["document"]
                    if type(doc) == list:
                        neighbours = [" ".join(doc)]
                    else:
                        neighbours = [doc]
                elif "sub-paragraphs" in instance:
                    neighbours = ["title: , source: " + instance["sub-paragraphs"]]
                elif fix_newsqa and "sub_paragraph" in instance:
                    neighbours = ["title: , source: " + instance["sub_paragraph"]]
                else:
                    neighbours = ["title: , source: "]

        if inference_only:
            data.append((question, None, neighbours))
        else:
            if True:
                if "answers" in instance:
                    answers = instance["answers"]
                elif "answer" in instance:
                    if type(instance["answer"]) is str:
                        answers = [instance["answer"]]
                    elif type(instance["answer"]) is list:
                        answers = instance["answer"]
                    else:
                        answers = [str(instance["answer"])]
                else:
                    raise ValueError("need to have answer or answers")
            if len(answers) < 1:
                continue
                # answers = ["This question cannot be answered based on the given information."]
            else:
                ## only take answer 0
                if type(answers[0]) is dict:
                    answers = [answers[0]["text"].strip()]
                elif type(answers[0]) is str:
                    answers = [answers[0]]
                else:
                    raise ValueError("unsupported type for answer(s)")

            for answer in answers:
                answer = format_answer(answer)
                data.append((question, answer, neighbours))
    
    return data


def reformat_prompt_v2(query, neighbours, dataset_name, ft_neighbours, \
    max_output_len, tokenizer, max_seq_length):

    system = "System: This is a chat between a user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions based on the context. The assistant should also indicate when the answer cannot be found in the context.\n\n"

    if dataset_name in ["oasst", "tuluv2", "tuluv2official", "quiet_cockatoo", "quiet-cockatoo_commercial", "primitive-stingray16k"]:
        if dataset_name == "tuluv2official":
            all_input = query
        else:
            all_input = system + query
            
        input_tokens = tokenizer.encode(all_input)
        return input_tokens

    short_span_with_context = ["drop", "NarrativeQA", "NarrativeQAretrieval", "QASC", "Quoref", "ROPES", "squad1.1", "squad2.0", "newsqa", "nq", "BioASQ", "DuoRC_ParaphraseRC", "TextbookQA", "WikiTableQuestions", "HybridQA", "hotpotqa", "wikimqa", "kilt_nq_short", "kilt_tqa_short", "kilt_hotpotqa_short", "nqtables", "qasper", "narrative_qa", "quality", "musique", "hotpotqa", "multifieldqa_en", "longbook_qa_eng", "kv_retrieval", "math_find", "passkey", "number_string", "code_debug", "code_run", "math_calc", "longdialogue_qa_eng", "longbook_qa_eng_gpt4_same", "longdialogue_qa_eng_gpt4_same" ]
    yes_no_without_context = ["boolq", "multirc"]
    multichoices = ["race", "longbook_choice_eng", "longbook_choice_eng_gpt4_same"]
    # multi-turn qa datasets
    formatted_dataset_name = ["convqa", "convqav2", "chatgptgen", "chatgptgennoanswer", "chatgptgennoanswerv2", "doc2dial", "doc2dialv2", "doc2dial_dragon", "quac", "quacv2", "quac_dragon", "qrecc", "qrecc_dragon", "sharc", "nvolvemultiturn1300", "nvolvemultiturn1700", "nvolvemultiturnfiltered5k", "nvolvemultiturnfiltered5knoanswer", "nvolvemultiturnfiltered7k", "nvolvemultiturnfiltered7knoanswer", "nvolvemultiturnfiltered7knoanswer1k", "nvolvemultiturnfiltered7knoanswer2k", "nvolvemultiturnfiltered7knoanswer3k",
    "nvolvemultiturnfiltered7knoanswerlonghistory", 
    "nvolvemultiturnfiltered7knoanswerlonghistorydiscont", "nvolvemultiturnfiltered7knoanswerlonghistorydiscontfixv1", "nvolvemultiturnfiltered7knoanswerlonghistorydiscontfixv2", "scalecqav1", "scalecqanoanswer", "extracqanoanswerv1", "instructv1", "instructv2", "instructv3", "instructtablegeneral", "instructtablegeneralv2", "instructtableunansreasongeneral", "doqa_cooking", "doqa_movies", "doqa_travel", "hybriddial", "hybriddial_general", "hybriddialunanswerable", "hybriddialunanswerablemixed", "hybriddialunanswerablegeneral", "hybriddialunanswerablegeneralv2", "inscit", "inscit_dragon", "convfinqalong", "convfinqalonggeneral", "convfinqalongunanswerable", "convfinqalongunanswerablegeneral", "convfinqalongunanswerablegeneralv2", "convfinqalongunanswerablewithreasongeneral", "cornercases", "cornercasesv2", "convfinqa_general_long_answer"]

    formatted_dataset_name_short = ["coqa"]
    formatted_dataset_name_short_and_long = ["sqa", "sqa_general", "topiocqa", "topiocqa_dragon"]
    formatted_dataset_name_entity = ["sqa_general_long_answer"]
    singleturn_dataset_name_short_and_long = ["tatqamultispan", "llmware", "tatqamultispangeneral"]
    singleturn_dataset_name_long = ["kilt_nq", "kilt_tqa", "kilt_hotpotqa", "kilt_hotpotqa_rerank"]
    singleturn_dataset_entity = ["tatqamultispanv2general"]

    math_program_with_context = ["finqa", "finqav2"]
    math_program_with_context_v2 = ['tatqav2']
    math_program_with_context_v3 = ['tatqav3', 'tatqageneral']
    math_program_multiturn = ["convfinqa", "convfinqav2"]
    math_program_multiturn_v2 = ["convfinqav3", "convfinqa_general"]

    user_template = ""

    if dataset_name in formatted_dataset_name:
        # dialogue_turn = query

        ## adding this instruction to multi-turn
        tmp_list = query.split("User:", 1)  # split will stop at the first "User:"
        dialogue_turn = "User: Please give a full and complete answer for the question." + tmp_list[1]
    
    elif dataset_name in formatted_dataset_name_short_and_long:

        tmp_list = query.split("User:")
        tmp_list = tmp_list[1:]

        dialogue_turn = ""
        if len(tmp_list) > 1:
            for item in tmp_list[:-1]:
                dialogue_turn += "User:" + item
        dialogue_turn += "User: Answer the following question with a short span, or a full and complete answer." + tmp_list[-1]
    
    elif dataset_name in formatted_dataset_name_entity:

        tmp_list = query.split("User:")
        tmp_list = tmp_list[1:]

        dialogue_turn = ""
        if len(tmp_list) > 1:
            for item in tmp_list[:-1]:
                dialogue_turn += "User:" + item
        dialogue_turn += "User: Answer the following question with one or a list of items." + tmp_list[-1]
    
    elif dataset_name in formatted_dataset_name_short:
        tmp_list = query.split("User:")
        tmp_list = tmp_list[1:]

        dialogue_turn = ""
        if len(tmp_list) > 1:
            for item in tmp_list[:-1]:
                dialogue_turn += "User:" + item
        dialogue_turn += "User: Answer the following question with a short span. The answer needs to be just in a few words." + tmp_list[-1]

    elif dataset_name in math_program_multiturn:

        ## for training
        tmp_list = query.split("User:", 1)  # split will stop at the first "User:"
        dialogue_turn = "User: Answer the following question with a number from context or the math arithmetic (add, subtract, multiply, and divide)." + tmp_list[1]

    elif dataset_name in math_program_multiturn_v2:
        ## for evaluation
        tmp_list = query.split("User:")
        tmp_list = tmp_list[1:]
        dialogue_turn = ""
        if len(tmp_list) > 1:
            for item in tmp_list[:-1]:
                dialogue_turn += "User:" + item
        dialogue_turn += "User: Answer the following question with a number from context or the math arithmetic using +, -, *, or /." + tmp_list[-1]

    else:
        if dataset_name in short_span_with_context:
            user = "Answer the following question with a short span. The answer needs to be just in a few words. {}".format(query)
        elif dataset_name in yes_no_without_context:
            user = "Answer the following question with True or False. {}".format(query)
        elif dataset_name in multichoices:
            user = "Answer the following question by selecting one of the provided options. {}".format(query)
        elif dataset_name in math_program_with_context:

            ## for evaluation
            user = "Answer the following question with the math arithmetic using +, -, *, or /. {}".format(query)
        elif dataset_name in math_program_with_context_v2:
            ## for evaluation
            user = "Answer the following question with a short span or a number from context or the math arithmetic (add, subtract, multiply, and divide). {}".format(query)
        elif dataset_name in math_program_with_context_v3:
            ## for training
            user = "Answer the following question with a number from context or the math arithmetic using +, -, *, or /. {}".format(query)

        elif dataset_name in singleturn_dataset_name_short_and_long:
            user = "Answer the following question with a short span, or a full and complete answer. {}".format(query)

        elif dataset_name in singleturn_dataset_name_long:
            user = "Please give a full and complete answer for the question. {}".format(query)

        elif dataset_name in singleturn_dataset_entity:
            user = "Answer the following question with one or a list of items. {}".format(query)

        elif dataset_name == "qmsum":
            user = "Please summarize a full and complete answer for the following question. {}".format(query)
        elif dataset_name == "longbook_sum_eng" or dataset_name == "longbook_sum_eng_gpt4_same":
            user = "Summarize the book above with a long paragraph."

        else:
            # fetaqa/llmware_unanswerable goes to here by default
            user = "Please give a full and complete answer for the question. {}".format(query)

        if dataset_name in ["kilt_nq_short", "kilt_tqa_short", "kilt_hotpotqa_short"]:
            dialogue_format = "User: {}\n\nAssistant: The answer is "
        else:
            dialogue_format = "User: {}\n\nAssistant:"
        dialogue_turn = dialogue_format.format(user)

    if ft_neighbours > 0:

        ## normal ordering
        context = "\n\n".join(neighbours[0:ft_neighbours]) + "\n\n"

        context_tokens = tokenizer.encode(context)
        dialogue_tokens = tokenizer.encode(dialogue_turn)
        system_tokens = tokenizer.encode(system)

        if len(system_tokens) + len(dialogue_tokens) + len(context_tokens) + max_output_len > max_seq_length:
            context_tokens = context_tokens[:max_seq_length - max_output_len - len(dialogue_tokens) - len(system_tokens)]
            context = tokenizer.decode(context_tokens, clean_up_tokenization_spaces=False) + "\n"
        
        all_input = system + context + dialogue_turn

        input_tokens = tokenizer.encode(all_input)
    else:
        all_input = system + dialogue_turn

        input_tokens = tokenizer.encode(all_input)

    return  input_tokens


def get_chatqa2_input(data_list, eval_dataset, tokenizer, num_ctx, max_output_len, max_seq_length):
    
    ft_neighbours = num_ctx
    dataset_name = eval_dataset
    prompt_list = []

    for sample in data_list:
        query, _, neighbours = sample
        input_tokens = reformat_prompt_v2(query, neighbours, dataset_name.split(".")[0], ft_neighbours, \
            max_output_len, tokenizer, max_seq_length)
        raw_text = tokenizer.decode(input_tokens, clean_up_tokenization_spaces=False)
        assert raw_text.startswith("<|begin_of_text|>")
        raw_text = raw_text[17:]
        prompt_list.append(raw_text)
    return prompt_list