from datetime import datetime
from pathlib import Path

from datasets import load_dataset

from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.normalizers import NFC
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.processors import TemplateProcessing
from tokenizers.trainers import BpeTrainer

# ==========================================================
# AYARLAR
# ==========================================================

DATASET_NAME = "Uunan/turkish-synthetic-corpus"

VOCAB_SIZE = 65536

OUTPUT_DIR = Path("tokenizer")
OUTPUT_DIR.mkdir(exist_ok=True)

SPECIAL_TOKENS = [
    "<pad>",
    "<unk>",
    "<bos>",
    "<eos>",
]

# ==========================================================
# BAŞLANGIÇ
# ==========================================================

start = datetime.now()

print("=" * 70)
print("         TURKISH BPE TOKENIZER TRAINING")
print("=" * 70)

print(f"Dataset        : {DATASET_NAME}")
print(f"Vocabulary     : {VOCAB_SIZE:,}")

# ==========================================================
# DATASET
# ==========================================================

print("\n[1/5] Dataset yükleniyor...")

dataset = load_dataset(DATASET_NAME)

print(dataset)

# ==========================================================
# TOKENIZER
# ==========================================================

print("\n[2/5] Tokenizer oluşturuluyor...")

tokenizer = Tokenizer(
    BPE(
        unk_token="<unk>"
    )
)

tokenizer.normalizer = NFC()

tokenizer.pre_tokenizer = ByteLevel(
    add_prefix_space=False
)

tokenizer.decoder = ByteLevelDecoder()

trainer = BpeTrainer(
    vocab_size=VOCAB_SIZE,
    min_frequency=2,
    special_tokens=SPECIAL_TOKENS,
    initial_alphabet=ByteLevel.alphabet(),
    show_progress=True,
)

print("✔ Hazır")

# ==========================================================
# DATASET ITERATOR
# ==========================================================

def batch_iterator(batch_size=1000):

    train = dataset["train"]

    for i in range(0, len(train), batch_size):
        yield train[i:i + batch_size]["text"]

# ==========================================================
# TRAIN
# ==========================================================

print("\n[3/5] Eğitim başlıyor...\n")

tokenizer.train_from_iterator(
    iterator=batch_iterator(),
    trainer=trainer,
    length=len(dataset["train"])
)

print("✔ Eğitim tamamlandı.")

# ==========================================================
# BOS EOS
# ==========================================================

print("\n[4/5] BOS/EOS ekleniyor...")

bos = tokenizer.token_to_id("<bos>")
eos = tokenizer.token_to_id("<eos>")

tokenizer.post_processor = TemplateProcessing(
    single="<bos> $A <eos>",
    pair="<bos> $A <eos> $B:1 <eos>:1",
    special_tokens=[
        ("<bos>", bos),
        ("<eos>", eos),
    ],
)

print("✔ Tamam")

# ==========================================================
# SAVE
# ==========================================================

print("\n[5/5] Kaydediliyor...")

tokenizer.save(str(OUTPUT_DIR / "tokenizer.json"))

elapsed = datetime.now() - start

print("\n" + "=" * 70)
print("TOKENIZER HAZIR")
print("=" * 70)

print(f"Süre            : {elapsed}")
print(f"Vocabulary      : {tokenizer.get_vocab_size():,}")
print(f"Dosya           : {OUTPUT_DIR/'tokenizer.json'}")

# ==========================================================
# TEST
# ==========================================================

sample = "Türkiye yapay zekâ alanında büyük yatırımlar yapıyor."

encoded = tokenizer.encode(sample)

print("\nÖrnek Metin")
print(sample)

print("\nTokenlar")
print(encoded.tokens)

print("\nID'ler")
print(encoded.ids)

decoded = tokenizer.decode(encoded.ids)

print("\nDecode")
print(decoded)
print("Not: ByteLevel BPE tokenizer'larında token listesinde görülen")
print("'Ġ', 'Ã', 'Ä' gibi karakterler tokenların byte gösterimidir.")
print("Bu bir hata değildir. Önemli olan encode → decode sonucunun")
print("orijinal metinle birebir aynı olmasıdır.\n")

print("Encode → Decode Başarılı mı?")
print(sample == decoded)