import regex  # Βιβλιοθήκη για υποστήριξη Unicode regex
from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, trainers
from tokenizers.pre_tokenizers import PreTokenizer
from datasets import load_dataset

# Ορισμός custom regex για κώδικα και πολυγλωσσικό κείμενο.
code_regex = r"""(?x:
    //.*?$|                # Σχόλια τύπου //
    /\*.*?\*/|             # Σχόλια τύπου /* */
    "(?:\\.|[^\\"])*"|     # Strings με διπλά εισαγωγικά
    '(?:\\.|[^\\'])*'|     # Strings με μονά εισαγωγικά
    \b(?:if|else|for|while|return|function)\b|  # Κλειδικά για κώδικα
    [<>]=?|\+\+|--|&&|\|\||[-+*/%=&|^~!]=?|  # Operators
    \d+\.?\d*|\.\d+|       # Αριθμοί
    [{}[\](),.;:]|         # Σύμβολα
    \p{L}+|\p{N}+|         # Unicode γράμματα/αριθμοί
    \s+|                   # Διαστήματα
    \S                     # Οποιοσδήποτε άλλος χαρακτήρας
)"""

def custom_pre_tokenizer(pretokenized_string):
    """
    Custom pre-tokenizer που χρησιμοποιεί το regex για να εξάγει tokens με offsets.
    """
    # Εφαρμόζουμε το regex για να βρούμε τα tokens
    tokens_with_offsets = [(m.group(), m.span()) for m in regex.finditer(code_regex, pretokenized_string.sequence)]
    pretokenized_string.split(lambda: tokens_with_offsets)

def train_tokenizer(iterator, vocab_size=32000, min_frequency=2):
    # Δημιουργία του Tokenizer με το μοντέλο BPE και ειδικό token για Unknown.
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    
    # Εφαρμογή normalization (NFC και αφαίρεση τόνων, αν χρειάζεται)
    tokenizer.normalizer = normalizers.Sequence([
        normalizers.NFC(),
        normalizers.StripAccents()  # Προαιρετικό: αφαιρεί τόνους
    ])
    
    # Ορισμός custom pre-tokenizer με χρήση της συνάρτησης που ορίσαμε
    tokenizer.pre_tokenizer = PreTokenizer.custom(custom_pre_tokenizer)
    
    # Επιπλέον χρήση του ByteLevel pre-tokenizer για συμβατότητα με το GPT-2
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
        tokenizer.pre_tokenizer,
        pre_tokenizers.ByteLevel(add_prefix_space=False)
    ])
    
    # Ορισμός του Trainer για το BPE
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=["<|endoftext|>", "<pad>", "<unk>", "<mask>"],
        continuing_subword_prefix="",
        show_progress=True
    )
    
    # Εκπαίδευση του tokenizer από iterator κειμένων
    tokenizer.train_from_iterator(iterator, trainer=trainer)
    
    # Ορισμός decoder ώστε να αντιστρέφεται σωστά η κωδικοποίηση
    tokenizer.decoder = decoders.ByteLevel()
    
    return tokenizer