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)