import torch

from transformers import ElectraTokenizerFast

class AnswerAwareTokenizer():
    def __init__(self, total_maxlen, bert_model='google/electra-base-discriminator'):
        self.total_maxlen = total_maxlen

        self.tok = ElectraTokenizerFast.from_pretrained(bert_model)

    def process(self, questions, passages, all_answers=None, mask=None):
        return TokenizationObject(self, questions, passages, all_answers, mask)

    def tensorize(self, questions, passages):
        query_lengths = self.tok(questions, padding='longest', return_tensors='pt').attention_mask.sum(-1)

        encoding = self.tok(questions, passages, padding='longest', truncation='longest_first',
                            return_tensors='pt', max_length=self.total_maxlen, add_special_tokens=True)

        return encoding, query_lengths

    def get_all_candidates(self, encoding, index):
        offsets, endpositions = self.all_word_positions(encoding, index)

        candidates = [(offset, endpos)
                      for idx, offset in enumerate(offsets)
                      for endpos in endpositions[idx:idx+10]]

        return candidates

    def all_word_positions(self, encoding, index):
        words = encoding.word_ids(index)
        offsets = [position
                   for position, (last_word_number, current_word_number) in enumerate(zip([-1] + words, words))
                   if last_word_number != current_word_number]

        endpositions = offsets[1:] + [len(words)]

        return offsets, endpositions

    def characters_to_tokens(self, text, answers, encoding, index, offset, endpos):
        # print(text, answers, encoding, index, offset, endpos)
        # endpos = endpos - 1

        for offset_ in range(offset, len(text)+1):
            tokens_offset = encoding.char_to_token(index, offset_)
            # print(f'tokens_offset = {tokens_offset}')
            if tokens_offset is not None:
                break

        for endpos_ in range(endpos, len(text)+1):
            tokens_endpos = encoding.char_to_token(index, endpos_)
            # print(f'tokens_endpos = {tokens_endpos}')
            if tokens_endpos is not None:
                break

        # None on whitespace!
        assert tokens_offset is not None, (text, answers, offset)
        # assert tokens_endpos is not None, (text, answers, endpos)
        tokens_endpos = tokens_endpos if tokens_endpos is not None else len(encoding.tokens(index))

        return tokens_offset, tokens_endpos

    def tokens_to_answer(self, encoding, index, text, tokens_offset, tokens_endpos):
        # print(encoding, index, text, tokens_offset, tokens_endpos, len(encoding.tokens(index)))

        char_offset = encoding.word_to_chars(index, encoding.token_to_word(index, tokens_offset)).start

        try:
            char_next_offset = encoding.word_to_chars(index, encoding.token_to_word(index, tokens_endpos)).start
            char_endpos = char_next_offset
        except:
            char_endpos = encoding.word_to_chars(index, encoding.token_to_word(index, tokens_endpos-1)).end

        assert char_offset is not None
        assert char_endpos is not None

        return text[char_offset:char_endpos].strip()


class TokenizationObject():
    def __init__(self, tokenizer: AnswerAwareTokenizer, questions, passages, answers=None, mask=None):
        assert type(questions) is list and type(passages) is list
        assert len(questions) in [1, len(passages)]

        if mask is None:
            mask = [True for _ in passages]

        self.mask = mask

        self.tok = tokenizer
        self.questions = questions if len(questions) == len(passages) else questions * len(passages)
        self.passages = passages
        self.answers = answers

        self.encoding, self.query_lengths = self._encode()
        self.passages_only_encoding, self.candidates, self.candidates_list = self._candidize()

        if answers is not None:
            self.gold_candidates = self.answers  # self._answerize()

    def _encode(self):
        return self.tok.tensorize(self.questions, self.passages)

    def _candidize(self):
        encoding = self.tok.tok(self.passages, add_special_tokens=False)

        all_candidates = [self.tok.get_all_candidates(encoding, index) for index in range(len(self.passages))]

        bsize, maxcands = len(self.passages), max(map(len, all_candidates))
        all_candidates = [cands + [(-1, -1)] * (maxcands - len(cands)) for cands in all_candidates]

        candidates = torch.tensor(all_candidates)
        assert candidates.size() == (bsize, maxcands, 2), (candidates.size(), (bsize, maxcands, 2), (self.questions, self.passages))

        candidates = candidates + self.query_lengths.unsqueeze(-1).unsqueeze(-1)

        return encoding, candidates, all_candidates