from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
from typing import List, Optional, Union
import os

MASK = "#"
MSA_PAD = "!"
UL_ALPHABET_PLUS = "ACDEFGHIKLMNPQRSTVWYBZXJOU-*#@!/[]{}"
MSA_AAS = "ACDEFGHIKLMNPQRSTVWYBZXJOU-"
GAP = "-"
START = "@"
STOP = "*"
SEP = "/"
END_AL = "]"
END_UL = "}"
START_AL = "["
START_UL = "{"

class ProteinTokenizer(PreTrainedTokenizer):

    def __init__(
        self,
        protein_alphabet: str = UL_ALPHABET_PLUS,
        model_max_length: int = 2048,
        pad_token=MSA_PAD,
        mask_token=MASK,
        all_aas=MSA_AAS,
        gap_token=GAP,
        bos_token=START,
        eos_token=STOP,
        sep_token=SEP,
        **kwargs
    ):
        """Character tokenizer for Hugging Face transformers.

        model_max_length (int): Model maximum sequence length.
        """
        self.alphabet = list("".join(protein_alphabet))
        self.all_aas = list("".join(all_aas))
        self.a_to_i = {u: i for i, u in enumerate(self.alphabet)}
        self.i_to_a = {i: u for i, u in enumerate(self.alphabet)}
        self.gap_token = gap_token

        
        bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
        eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
        sep_token = AddedToken(sep_token, lstrip=False, rstrip=False) if isinstance(sep_token, str) else sep_token
        mask_token = AddedToken(mask_token, lstrip=False, rstrip=False) if isinstance(mask_token, str) else mask_token 
        pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
        gap_token = AddedToken(gap_token, lstrip=False, rstrip=False) if isinstance(gap_token, str) else gap_token

        super().__init__(
            pad_token=pad_token,
            mask_token=mask_token,
            eos_token=eos_token,
            bos_token=bos_token,
            sep_token=sep_token,
            model_max_length=model_max_length,
            **kwargs
        )

    @property
    def vocab_size(self):
        return len(self.alphabet)
    
    @property
    def gap_token_id(self):
        return self.convert_tokens_to_ids(self.gap_token)

    def get_vocab(self):
        return self.a_to_i

    def _tokenize(self, text: str) -> List[str]:
        return list(text)

    def _convert_token_to_id(self, token) -> int:
        return self.a_to_i[token]

    def _convert_id_to_token(self, index) -> str:
        return self.i_to_a[index]

    def convert_tokens_to_string(self, tokens):
        return "".join(tokens)

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        result = token_ids_0
        if token_ids_1 is not None:
            raise NotImplementedError("This tokenizer does not support two sequences")
        return result

    def get_special_tokens_mask(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
        already_has_special_tokens: bool = False,
    ) -> List[int]:
        if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_0=token_ids_0,
                token_ids_1=token_ids_1,
                already_has_special_tokens=True,
            )

        result = [0] * len(token_ids_0)
        if token_ids_1 is not None:
            raise NotImplementedError("This tokenizer does not support two sequences")

        return result

    def create_token_type_ids_from_sequences(
        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
    ) -> List[int]:
        """
        Identifies the type of token. 0 for the first sentence, 1 for the second sentence if it exists
        """

        result = len(token_ids_0) * [0]

        if token_ids_1 is not None:
            raise NotImplementedError("This tokenizer does not support two sequences")
        return result

    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
        super().save_pretrained(save_directory, **kwargs)
    
    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None):
        return ()