import os
import re

from transformers import PreTrainedTokenizer


class GenoJEPATokenizer(PreTrainedTokenizer):
    def __init__(self, **kwargs):
        self.pad_token = "N"

        self.base_chars = ["A", "T", "C", "G"]
        self.special_tokens = [self.pad_token]

        self.vocab = {token: i for i, token in enumerate(self.base_chars + self.special_tokens)}
        self.ids_to_tokens = {i: token for token, i in self.vocab.items()}

        self.dna_pattern = re.compile(r"[ATCG]")
        self.pad_token_id = self._convert_token_to_id(self.pad_token)

        super().__init__(**kwargs)

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

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

    def _convert_token_to_id(self, token):
        return self.vocab.get(token, self.vocab[self.pad_token])

    def _convert_id_to_token(self, index):
        return self.ids_to_tokens.get(index, self.pad_token)

    def _tokenize(self, seq, **kwargs):
        tokens = []
        pos = 0
        while pos < len(seq):
            dna_match = self.dna_pattern.match(seq, pos)
            if dna_match:
                dna_seq = dna_match.group()
                tokens.append(dna_seq)
                pos = dna_match.end()
            else:
                tokens.append(self.pad_token)
                pos += 1

        return tokens

    def save_vocabulary(self, save_directory, filename_prefix=None):
        vocab_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.txt")
        with open(vocab_file, "w", encoding="utf-8") as writer:
            for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
                writer.write(token + "\n")

        return (vocab_file,)