# This file contains the inference code for loading and running the closed-book and open-book QA models
import os
import csv
import glob
import gzip
import string
import sys
from typing import List, Tuple, Dict
import re
import math
import collections

import numpy as np
import unicodedata
import torch
from torch import Tensor as T
from torch import nn

from DPR.dpr.models import init_biencoder_components
from DPR.dpr.options import setup_args_gpu, print_args, set_encoder_params_from_state 
from DPR.dpr.indexer.faiss_indexers import DenseIndexer, DenseFlatIndexer
from DPR.dpr.data.reader_data import ReaderSample, ReaderPassage, get_best_spans
from DPR.dpr.models import init_reader_components
from DPR.dpr.utils.data_utils import Tensorizer
from DPR.dpr.utils.model_utils import load_states_from_checkpoint, get_model_obj

from transformers import GPT2LMHeadModel, GPT2Tokenizer, T5ForConditionalGeneration, AutoTokenizer
from segment_fill import segment_fill

sys.path.append("../.")
sys.path.append("./")
sys.path.append("./DPR/")

abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)

SEGMENTER_CACHE = {}
RERANKER_CACHE = {}

def setup_t5_reranker(process_id):
    tokenizer = AutoTokenizer.from_pretrained('google/byt5-small')
    model = T5ForConditionalGeneration.from_pretrained('closedbook/byte-t5-reranker/newest-test-9000')
    model.eval().to('cuda:'+str(process_id % torch.cuda.device_count()))
    return model, tokenizer

def t5_reranker_score_with_clue(model, tokenizer, clues, possibly_ungrammatical_fills):
    global RERANKER_CACHE
    results = []
    for clue, possibly_ungrammatical_fill in zip(clues, possibly_ungrammatical_fills):
        if not possibly_ungrammatical_fill.islower():
            possibly_ungrammatical_fill = possibly_ungrammatical_fill.lower()
        clue = preprocess_clue_fn(clue)
        if clue[-3:] == '. .':
            clue = clue[:-3]
        elif clue[-3:] == ' ..':
            clue = clue[:-3]
        elif clue[-2:] == '..':
            clue = clue[:-2]
        elif clue[-1] == '.':
            clue = clue[:-1]

        if clue + possibly_ungrammatical_fill in RERANKER_CACHE:
            results.append(RERANKER_CACHE[clue + possibly_ungrammatical_fill])
            continue
        else:
            with torch.inference_mode():
                inputs = tokenizer(['Q: ' + clue], return_tensors='pt')['input_ids'].to(model.device)
                labels = tokenizer([possibly_ungrammatical_fill], return_tensors='pt')['input_ids'].to(model.device)
                loss = model(inputs, labels=labels)
                answer_length = labels.shape[1]
                logprob = -loss[0].item() * answer_length
                results.append(logprob)
                RERANKER_CACHE[clue + possibly_ungrammatical_fill] = logprob

    return results

def chunks(lst, n=32):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i:i + n]

def batched_t5_reranker_score_with_clue(model, tokenizer, clues, possibly_ungrammatical_fills):
    global RERANKER_CACHE
    results = []
    clues_to_run = []
    fills_to_run = []
    with torch.inference_mode():
        for index, (clue, possibly_ungrammatical_fill) in enumerate(zip(clues, possibly_ungrammatical_fills)):
            if not possibly_ungrammatical_fill.islower():
                possibly_ungrammatical_fills[index]= possibly_ungrammatical_fill.lower()
            
            clue = preprocess_clue_fn(clue)
            if clue[-3:] == '. .':
                clue = clue[:-3]
            elif clue[-3:] == ' ..':
                clue = clue[:-3]
            elif clue[-2:] == '..':
                clue = clue[:-2]
            elif clue[-1] == '.':
                clue = clue[:-1]
            clues[index] = clue
            if clues[index] + possibly_ungrammatical_fills[index] in RERANKER_CACHE:
                results.append(RERANKER_CACHE[clues[index] + possibly_ungrammatical_fills[index]])
            else:
                results.append(1000.0)
                clues_to_run.append(clues[index])
                fills_to_run.append(possibly_ungrammatical_fills[index])

        if len(clues_to_run) > 0:
            for clues_chunk, fills_chunk in zip(list(chunks(clues_to_run)), list(chunks(fills_to_run))): 
                inputs = tokenizer(['Q: ' + clue for clue in clues_chunk], return_tensors='pt', padding='longest')['input_ids'].to(model.device)
                labels = tokenizer(fills_chunk, return_tensors='pt', padding='longest')['input_ids'].to(model.device)
                output = model(inputs, labels=labels)
                logits = output['logits']
                loss = nn.CrossEntropyLoss(reduction='none', ignore_index=0)(logits.view(-1, logits.size(-1)), labels.view(-1))
                loss = loss.view(inputs.shape[0], -1)
                loss[labels == 0] = 0
                loss = -1 * loss.sum(-1)
                loss = loss.detach().cpu().numpy().tolist()

                curr_index = 0
                for indx, result in enumerate(results):
                    if curr_index >= len(clues_chunk):
                        break
                    if result == 1000.0:
                        results[indx] = loss[curr_index]
                        RERANKER_CACHE[clues_chunk[curr_index] + fills_chunk[curr_index]] = loss[curr_index]
                        curr_index += 1

    return results

def setup_lm_reranker(process_id):
    model = GPT2LMHeadModel.from_pretrained("lm_reranker/large-checkpoint-330000")#checkpoint-494000/")
    model.eval().to('cuda:'+str(process_id % torch.cuda.device_count()))
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    return model, tokenizer

def lm_reranker_score_with_clue(model, tokenizer, clues, fills):
    global SEGMENTER_CACHE
    global RERANKER_CACHE
    results = []
    with torch.inference_mode():
        uncached_fills = [fill for fill in fills if fill not in SEGMENTER_CACHE]
        if len(uncached_fills) > 0:
            uncached_segmented_sentences = segment_fill(uncached_fills, 'segmenter/checkpoints/checkpoint-275000')
            for fill, uss in zip(uncached_fills, uncached_segmented_sentences):
                SEGMENTER_CACHE[fill] = uss
        segmented_sentences = [SEGMENTER_CACHE[fill] for fill in fills]
        for c, ss in zip(clues, segmented_sentences):
            c = preprocess_clue_fn(c)
            ss = ss.lower()
            if (c, ss) in RERANKER_CACHE:
                results.append(RERANKER_CACHE[(c, ss)])
            else:
                tensor_input = tokenizer.encode('Question: ' + c + '? Answer: ' + ss, return_tensors='pt').to(model.device)
                target_tensor = tensor_input.clone()
                seen_answer_token = False
                for idx, token in enumerate(target_tensor[0]):
                    if target_tensor[0][idx] == 25 and seen_answer_token:
                        target_tensor[0][idx] = -100
                        break
                    elif target_tensor[0][idx] == 23998:
                        seen_answer_token = True
                    target_tensor[0][idx] = -100
                loss=model(tensor_input, labels=target_tensor)
                answer_length = sum(target_tensor[0]!=-100).item()
                final_score = -loss[0].item() * answer_length
                results.append(final_score) # TODO, consider other scoring functions than log likelihood
                RERANKER_CACHE[(c, ss)] = final_score
    return results

def preprocess_clue_fn(clue):
    clue = str(clue)

    # https://stackoverflow.com/questions/517923/what-is-the-best-way-to-remove-accents-normalize-in-a-python-unicode-string
    clue = ''.join(c for c in unicodedata.normalize('NFD', clue) if unicodedata.category(c) != 'Mn')

    clue = re.sub("\x17|\x18|\x93|\x94|“|”|''|\"\"", "\"", clue)
    clue = re.sub("\x85|…", "...", clue)
    clue = re.sub("\x91|\x92|‘|’", "'", clue)

    clue = re.sub("‚", ",", clue)
    clue = re.sub("—|–", "-", clue)
    clue = re.sub("¢", " cents", clue)
    clue = re.sub("¿|¡|^;|\{|\}", "", clue)
    clue = re.sub("÷", "division", clue)
    clue = re.sub("°", " degrees", clue)

    euro = re.search("^£[0-9]+(,*[0-9]*){0,}| £[0-9]+(,*[0-9]*){0,}", clue)
    if euro:
        num = clue[:euro.end()]
        rest_clue = clue[euro.end():]
        clue = num + " Euros" + rest_clue
        clue = re.sub(", Euros", " Euros", clue)
        clue = re.sub("Euros [Mm]illion", "million Euros", clue)
        clue = re.sub("Euros [Bb]illion", "billion Euros", clue)
        clue = re.sub("Euros[Kk]", "K Euros", clue)
        clue = re.sub(" K Euros", "K Euros", clue)
        clue = re.sub("£", "", clue)

    clue = re.sub(" *\(\d{1,},*\)$| *\(\d{1,},* \d{1,}\)$", "", clue)

    clue = re.sub("&amp;", "&", clue)
    clue = re.sub("&lt;", "<", clue)
    clue = re.sub("&gt;", ">", clue)

    clue = re.sub("e\.g\.|for ex\.", "for example", clue)
    clue = re.sub(": [Aa]bbreviat\.|: [Aa]bbrev\.|: [Aa]bbrv\.|: [Aa]bbrv|: [Aa]bbr\.|: [Aa]bbr", " abbreviation", clue)
    clue = re.sub("abbr\.|abbrv\.", "abbreviation", clue)
    clue = re.sub("Abbr\.|Abbrv\.", "Abbreviation", clue)
    clue = re.sub("\(anag\.\)|\(anag\)", "(anagram)", clue)
    clue = re.sub("org\.", "organization", clue)
    clue = re.sub("Org\.", "Organization", clue)
    clue = re.sub("Grp\.|Gp\.", "Group", clue)
    clue = re.sub("grp\.|gp\.", "group", clue)
    clue = re.sub(": Sp\.", " (Spanish)", clue)
    clue = re.sub("\(Sp\.\)|Sp\.", "(Spanish)", clue)
    clue = re.sub("Ave\.", "Avenue", clue)
    clue = re.sub("Sch\.", "School", clue)
    clue = re.sub("sch\.", "school", clue)
    clue = re.sub("Agcy\.", "Agency", clue)
    clue = re.sub("agcy\.", "agency", clue)
    clue = re.sub("Co\.", "Company", clue)
    clue = re.sub("co\.", "company", clue)
    clue = re.sub("No\.", "Number", clue)
    clue = re.sub("no\.", "number", clue)
    clue = re.sub(": [Vv]ar\.", " variable", clue)
    clue = re.sub("Subj\.", "Subject", clue)
    clue = re.sub("subj\.", "subject", clue)
    clue = re.sub("Subjs\.", "Subjects", clue)
    clue = re.sub("subjs\.", "subjects", clue)

    theme_clue = re.search("^.+\|[A-Z]{1,}", clue)
    if theme_clue:
        clue = re.sub("\|", " | ", clue)

    if "Partner of" in clue:
        clue = re.sub("Partner of", "", clue)
        clue = clue + " and ___"

    link = re.search("^.+-.+ [Ll]ink$", clue)
    if link:
        no_link = re.search("^.+-.+ ", clue)
        x_y = clue[no_link.start():no_link.end() - 1]
        x_y_lst = x_y.split("-")
        clue = x_y_lst[0] + " ___ " + x_y_lst[1]

    follower = re.search("^.+ [Ff]ollower$", clue)
    if follower:
        no_follower = re.search("^.+ ", clue)
        x = clue[:no_follower.end() - 1]
        clue = x + " ___"

    preceder = re.search("^.+ [Pp]receder$", clue)
    if preceder:
        no_preceder = re.search("^.+ ", clue)
        x = clue[:no_preceder.end() - 1]
        clue = "___ " + x

    if re.search("--[^A-Za-z]|--$", clue):
        clue = re.sub("--", "__", clue)
    if not re.search("_-[A-Za-z]|_-$", clue):
        clue = re.sub("_-", "__", clue)

    clue = re.sub("_{2,}", "___", clue)

    clue = re.sub("\?$", " (wordplay)", clue)

    nonverbal = re.search("\[[^0-9]+,* *[^0-9]*\]", clue)
    if nonverbal:
        clue = re.sub("\[|\]", "", clue)
        clue = clue + " (nonverbal)"

    if clue[:4] == "\"\"\" " and clue[-4:] == " \"\"\"":
        clue = "\"" + clue[4:-4] + "\""
    if clue[:4] == "''' " and clue[-4:] == " '''":
        clue = "'" + clue[4:-4] + "'"
    if clue[:3] == "\"\"\"" and clue[-3:] == "\"\"\"":
        clue = "\"" + clue[3:-3] + "\""
    if clue[:3] == "'''" and clue[-3:] == "'''":
        clue = "'" + clue[3:-3] + "'"

    return clue

def setup_closedbook(process_id):
    dpr = DPRForCrossword(
        "closedbook/dpr_biencoder.12.50163",
        "closedbook/wordlist.tsv",
        "closedbook/embeddings/embeddings.json_*",
        retrievalmodel=False,
        process_id=process_id
    )

    # run some clues through to verify it runs and to warmup the devices
    clues = ['Test clue here', 'Test clue here #2', 'Clue with weird characters asdfghjkl;,./,/.`3`~!@#$%^&*()_+']
    outputs = answer_clues(dpr, clues, 200, output_strings=True)
    return dpr

def answer_clues(dpr, clues, max_answers, output_strings=False):
    clues = [preprocess_clue_fn(c.rstrip()) for c in clues]
    outputs = dpr.answer_clues_closedbook(clues, max_answers, output_strings=output_strings)
    return outputs

class DenseRetriever(object):
    """
    Does passage retrieving over the provided index and question encoder
    """
    def __init__(
        self,
        question_encoder: nn.Module,
        batch_size: int,
        tensorizer: Tensorizer,
        index: DenseIndexer,
        device=None,
    ):
        self.question_encoder = question_encoder
        self.batch_size = batch_size
        self.tensorizer = tensorizer
        self.index = index
        self.device = device

    def generate_question_vectors(self, questions: List[str]) -> T:
        n = len(questions)
        bsz = self.batch_size
        query_vectors = []
        self.question_encoder.eval()

        with torch.no_grad():
            for j, batch_start in enumerate(range(0, n, bsz)):
                batch_token_tensors = [
                    self.tensorizer.text_to_tensor(q)
                    for q in questions[batch_start : batch_start + bsz]
                ]

                q_ids_batch = torch.stack(batch_token_tensors, dim=0).to(self.device)
                q_seg_batch = torch.zeros_like(q_ids_batch).to(self.device)
                q_attn_mask = self.tensorizer.get_attn_mask(q_ids_batch)
                _, out, _ = self.question_encoder(q_ids_batch, q_seg_batch, q_attn_mask)

                query_vectors.extend(out.cpu().split(1, dim=0))
        query_tensor = torch.cat(query_vectors, dim=0)
        assert query_tensor.size(0) == len(questions)
        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
        :param query_vectors:
        :param top_docs:
        :return:
        """
        results = self.index.search_knn(query_vectors, top_docs)
        return results

class FakeRetrieverArgs:
    """Used to surpress the existing argparse inside DPR so we can have our own argparse"""
    def __init__(self):
        self.do_lower_case = False
        self.pretrained_model_cfg = None
        self.encoder_model_type = None
        self.model_file = None
        self.projection_dim = 0
        self.sequence_length = 512
        self.do_fill_lower_case = False
        self.desegment_valid_fill = False
        self.no_cuda = False
        self.local_rank = -1
        self.fp16 = False
        self.fp16_opt_level = "O1"


class DPRForCrossword(object):
    """Closedbook model for Crossword clue answering"""

    def __init__(
        self,
        model_file,
        ctx_file,
        encoded_ctx_file,
        batch_size=6000,
        retrievalmodel=False,
        process_id=0
    ):
        self.retrievalmodel = retrievalmodel  # am I a wikipedia retrieval model or a closed-book model
        args = FakeRetrieverArgs()
        args.model_file = model_file
        args.ctx_file = ctx_file
        args.encoded_ctx_file = encoded_ctx_file
        args.batch_size = batch_size
        self.device = torch.device("cuda:"+str(process_id%torch.cuda.device_count()))

        setup_args_gpu(args)
        print_args(args)

        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)

        question_encoder = encoder.question_model
        question_encoder = question_encoder.to(self.device)
        question_encoder.eval()

        # load weights from the model file
        model_to_load = get_model_obj(question_encoder)

        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()

        index = DenseFlatIndexer(vector_size, 50000)

        self.retriever = DenseRetriever(
            question_encoder,
            args.batch_size,
            tensorizer,
            index,
            self.device,
        )

        # index all passages
        ctx_files_pattern = args.encoded_ctx_file
        input_paths = glob.glob(ctx_files_pattern)
        self.retriever.index.index_data(input_paths)

        self.all_passages = self.load_passages(args.ctx_file)
        self.fill2id = {}
        for key in self.all_passages.keys():
            self.fill2id[
                "".join(
                    [
                        letter
                        for letter in self.all_passages[key][1].upper()
                        if letter in string.ascii_uppercase
                    ]
                )
            ] = key

        # might as well uppercase and remove non-alphas from the fills before we start to save time later
        if not retrievalmodel:
            temp = {}
            for my_id in self.all_passages.keys():
                temp[my_id] = "".join([c.upper() for c in self.all_passages[my_id][1] if c.upper() in string.ascii_uppercase])
            self.len_all_passages = len(list(self.all_passages.values()))
            self.all_passages = temp
            

    @staticmethod
    def load_passages(ctx_file: str) -> Dict[object, Tuple[str, str]]:
        docs = {}
        if ctx_file.endswith(".gz"):
            with gzip.open(ctx_file, "rt") as tsvfile:
                reader = csv.reader(
                    tsvfile,
                    delimiter="\t",
                )
                # file format: doc_id, doc_text, title
                for row in reader:
                    if row[0] != "id":
                        docs[row[0]] = (row[1], row[2])
        else:
            with open(ctx_file) as tsvfile:
                reader = csv.reader(
                    tsvfile,
                    delimiter="\t",
                )
                # file format: doc_id, doc_text, title
                for row in reader:
                    if row[0] != "id":
                        docs[row[0]] = (row[1], row[2])
        return docs

    def answer_clues_closedbook(self, questions, max_answers, output_strings=False):
        # assumes clues are preprocessed
        assert self.retrievalmodel == False
        questions_tensor = self.retriever.generate_question_vectors(questions)

        if max_answers > self.len_all_passages:
            max_answers = self.len_all_passages

        # get top k results
        top_ids_and_scores = self.retriever.get_top_docs(questions_tensor.numpy(), max_answers)
        
        if not output_strings:
            return top_ids_and_scores
        else:
            # get the string forms
            all_answers = []
            all_scores = []
            for ans in top_ids_and_scores: 
                all_answers.append(list(map(self.all_passages.get, ans[0])))
                all_scores.append(ans[1])
            return all_answers, all_scores
