from typing import Union, List

from transformers import AutoTokenizer
import torch


class HFPTTokenizer(object):
    def __init__(self, pt_name=None):

        self.pt_name = pt_name
        self.added_sep_token = 0
        self.added_cls_token = 0
        self.enable_add_tokens = False
        self.gpt_special_case = ((not self.enable_add_tokens) and ('gpt' in self.pt_name))

        if (pt_name is None):
            self.tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(pt_name)

        # Adding tokens to GPT causing NaN training loss.
        # Disable for now until further investigation.
        if (self.enable_add_tokens):
            if (self.tokenizer.sep_token is None):
                self.tokenizer.add_special_tokens({'sep_token': '<SEP>'})
                self.added_sep_token = 1

            if (self.tokenizer.cls_token is None):
                self.tokenizer.add_special_tokens({'cls_token': '<CLS>'})
                self.added_cls_token = 1

        if (self.gpt_special_case):
            self.tokenizer.pad_token = self.tokenizer.eos_token
            self.tokenizer.sep_token = self.tokenizer.eos_token

    def get_eot_token(self):
        return self.tokenizer.encode(self.tokenizer.sep_token, add_special_tokens=False)[0]

    def get_sot_token(self):
        return self.tokenizer.encode(self.tokenizer.cls_token, add_special_tokens=False)[0]

    def get_eot_token_list(self):
        return self.tokenizer.encode(self.tokenizer.sep_token, add_special_tokens=False)

    def get_sot_token_list(self):
        return self.tokenizer.encode(self.tokenizer.cls_token, add_special_tokens=False)

    def get_tokenizer_obj(self):
        return self.tokenizer

    # Language model needs to know if new tokens
    # were added to the dictionary.
    def check_added_tokens(self):
        return self.added_sep_token + self.added_cls_token

    def tokenize(self, texts: Union[str, List[str]], context_length: int = 77):
        if isinstance(texts, str):
            texts = [texts]

        padding = 'max_length'

        seqstart = []
        seqtok = []
        seqend = []

        max_length = context_length

        if (self.added_cls_token > 0):
            seqstart = self.get_sot_token_list()
            max_length = max_length - 1

        if (self.added_sep_token > 0):
            seqend = self.get_eot_token_list()
            max_length = max_length - 1

        tokens = self.tokenizer(
            texts, padding=padding,
            truncation=True,
            max_length=max_length
        )['input_ids']

        for i in range(len(tokens)):
            tokens[i] = seqstart + tokens[i] + seqend

        if (self.gpt_special_case):
            for i in range(len(tokens)):
                tokens[i][-1] = self.get_eot_token()

        # print(str(tokens))

        result = torch.Tensor(tokens).type(torch.LongTensor)

        return result

    def get_vocab_size(self):
        return self.tokenizer.vocab_size

    def __call__(self, texts: Union[str, List[str]], context_length: int = 77):
        return self.tokenize(texts, context_length)