from typing import Tuple
from transformers import DistilBertForQuestionAnswering, DistilBertTokenizer
import numpy as np
import torch
from nboost.models.qa import QAModel


def _is_whitespace(c):
    if c in " \t\r\n" or ord(c) == 0x202F:
        return True
    return False


class PtDistilBertQAModel(QAModel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = DistilBertForQuestionAnswering.from_pretrained(self.model_dir)
        self.tokenizer = DistilBertTokenizer.from_pretrained(self.model_dir)

    def get_answer(self, query: str, choice: str) -> Tuple[str, Tuple[int, int, int], float]:
        """Return (answer, (candidate, start_pos, end_pos))"""
        doc_tokens = []
        char_to_word_offset = []
        prev_is_whitespace = True

        # Split on whitespace so that different tokens may be attributed to
        # their original position.
        for c in choice:
            if _is_whitespace(c):
                prev_is_whitespace = True
            else:
                if prev_is_whitespace:
                    doc_tokens.append(c)
                else:
                    doc_tokens[-1] += c
                prev_is_whitespace = False
            char_to_word_offset.append(len(doc_tokens) - 1)

        tok_to_orig_index = []
        all_doc_tokens = []
        for (i, token) in enumerate(doc_tokens):
            sub_tokens = self.tokenizer.tokenize(token)
            for sub_token in sub_tokens:
                tok_to_orig_index.append(i)
                all_doc_tokens.append(sub_token)

        truncated_query = self.tokenizer.encode(query,
                                                add_special_tokens=False,
                                                max_length=self.max_query_length)

        encoded_dict = self.tokenizer.encode_plus(
            truncated_query,
            all_doc_tokens,
            max_length=self.max_seq_len,
            return_tensors='pt'
        )

        self.model.eval()
        with torch.no_grad():
            start_logits, end_logits = self.model(
                input_ids=encoded_dict['input_ids'])
            # add +2 for [CLS] and [SEP], and cut out last [SEP]
            start_logits = start_logits[0][len(truncated_query) + 2:-1]
            end_logits = end_logits[0][len(truncated_query) + 2:-1]

        assert len(end_logits) == len(tok_to_orig_index) or len(end_logits) == \
               self.max_seq_len - len(truncated_query) - 3

        max_score = start_logits[0] + end_logits[-1]
        start_tok = 0
        end_tok = len(end_logits) - 1
        for i, start_logit in enumerate(start_logits[:-1]):
            end_logit_pos = np.argmax(end_logits[i+1:]) + i + 1
            score = start_logit + end_logits[end_logit_pos]
            if score > max_score:
                max_score = score
                start_tok = i
                end_tok = end_logit_pos

        answer = ' '.join(doc_tokens[
                          tok_to_orig_index[start_tok]
                          :tok_to_orig_index[end_tok] + 1])
        start_char_offset = char_to_word_offset.index(
            tok_to_orig_index[start_tok])
        end_char_offset = char_to_word_offset.index(
            tok_to_orig_index[min(end_tok+1, len(tok_to_orig_index)-1)]) - 1
        return answer, (start_char_offset, end_char_offset, 0), max_score
