import json
import os
from transformers import PreTrainedTokenizer

VOCAB = {"<pad>": 0, "<mask>": 1, "A": 2, "U": 3, "G": 4, "C": 5}


class RNABertTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "vocab.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file=None,
        pad_token="<pad>",
        mask_token="<mask>",
        unk_token="<pad>",
        **kwargs,
    ):
        self._vocab = dict(VOCAB)
        if vocab_file and os.path.isfile(vocab_file):
            with open(vocab_file) as f:
                self._vocab = json.load(f)
        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
        super().__init__(
            pad_token=pad_token,
            mask_token=mask_token,
            unk_token=unk_token,
            **kwargs,
        )

    @property
    def vocab_size(self):
        return len(self._vocab)

    def get_vocab(self):
        return dict(self._vocab)

    def _tokenize(self, text):
        return list(text.upper().replace("T", "U"))

    def _convert_token_to_id(self, token):
        return self._vocab.get(token, 0)

    def _convert_id_to_token(self, index):
        return self._ids_to_tokens.get(index, "<pad>")

    def save_vocabulary(self, save_directory, filename_prefix=None):
        os.makedirs(save_directory, exist_ok=True)
        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
        path = os.path.join(save_directory, fname)
        with open(path, "w") as f:
            json.dump(self._vocab, f, indent=2)
        return (path,)

    @property
    def cls_token_id(self):
        return self.pad_token_id

    @property
    def eos_token_id(self):
        return self.pad_token_id

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        if token_ids_1 is None:
            return token_ids_0
        return token_ids_0 + token_ids_1

    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
        if already_has_special_tokens:
            return super().get_special_tokens_mask(token_ids_0, token_ids_1, True)
        return [0] * len(token_ids_0) + ([0] * len(token_ids_1) if token_ids_1 else [])

    def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
        if token_ids_1 is None:
            return [0] * len(token_ids_0)
        return [0] * len(token_ids_0) + [0] * len(token_ids_1)