from transformers import PreTrainedTokenizer
import numpy as np
import re

def replace_halogen(smiles):
    return smiles.replace('Cl', 'L').replace('Br', 'R')

def restore_halogen(smiles):
    return smiles.replace('L', 'Cl').replace('R', 'Br')

class CustomSMILESTokenizer(PreTrainedTokenizer):
    def __init__(self, vocab_file=None, max_length=140, **kwargs):
        self.special_tokens = ['[MASK]']
        self.additional_chars = []
        self.max_length = max_length

        if vocab_file:
            with open(vocab_file, 'r') as f:
                self.additional_chars = sorted(f.read().split())

        self.chars = self.special_tokens + self.additional_chars
        self.vocab = dict(zip(self.chars, range(len(self.chars))))
        self.inv_vocab = {v: k for k, v in self.vocab.items()}

        super().__init__(**kwargs)

    def _tokenize(self, text):
        regex = '(\[[^\[\]]{1,6}\])'
        text = replace_halogen(text)
        char_list = re.split(regex, text)
        tokens = []
        for char in char_list:
            if char.startswith('['):
                tokens.append(char)
            else:
                tokens.extend(list(char))
        return tokens

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

    def _convert_id_to_token(self, index):
        return self.inv_vocab.get(index, '[MASK]')

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

    def get_vocab(self):
        return self.vocab

    def save_vocabulary(self, save_directory, filename_prefix=None):
        file = save_directory + "/vocab.txt"
        with open(file, "w") as f:
            for token in self.additional_chars:
                f.write(token + "\n")
        return (file,)