"""
Train một BPE Tokenizer từ file output.csv
Chạy: python src/train_tokenizer.py
Kết quả: data/processed/bpe_tokenizer.json
"""
import os
import sys
import csv

sys.stdout.reconfigure(encoding='utf-8')

INPUT_CSV = 'data/raw/output.csv'
TOKENIZER_PATH = 'data/processed/bpe_tokenizer.json'

os.makedirs(os.path.dirname(TOKENIZER_PATH), exist_ok=True)


def main():
    from tokenizers import ByteLevelBPETokenizer

    if not os.path.exists(INPUT_CSV):
        print(f"Không tìm thấy {INPUT_CSV}")
        return

    # Thu thập tất cả câu hỏi + câu trả lời thành 1 file text tạm
    tmp_txt = 'data/processed/_tmp_corpus.txt'
    with open(INPUT_CSV, 'r', encoding='utf-8-sig') as fin, \
         open(tmp_txt, 'w', encoding='utf-8') as fout:
        reader = csv.reader(fin)
        header = next(reader, None)  # bỏ header
        for row in reader:
            if len(row) >= 3:
                fout.write(row[1].strip() + '\n')  # question
                fout.write(row[2].strip() + '\n')  # answer

    # Khởi tạo ByteLevel BPE tokenizer để giữ nguyên khoảng trắng
    tokenizer = ByteLevelBPETokenizer(lowercase=True)

    # Huấn luyện
    tokenizer.train(
        files=[tmp_txt],
        vocab_size=8000,
        min_frequency=2,
        special_tokens=["<pad>", "<sos>", "<eos>", "<unk>", "<sep>"],
        show_progress=True
    )
    
    tokenizer.save(TOKENIZER_PATH)

    # Dọn file tạm
    os.remove(tmp_txt)

    # Test thử
    print(f"\n{'='*55}")
    print(f"  BPE TOKENIZER TRAIN XONG!")
    print(f"{'='*55}")
    print(f"  Vocab size : {tokenizer.get_vocab_size()}")
    print(f"  Saved to   : {TOKENIZER_PATH}")
    print()

    test_sentences = [
        "Luật lũy đẳng là gì?",
        "Giải thích thuật toán Dijkstra.",
        "Đường đi Euler khác đường đi Hamilton chỗ nào?",
    ]
    for s in test_sentences:
        enc = tokenizer.encode(s)
        print(f"  Input : {s}")
        print(f"  Tokens: {enc.tokens}")
        print(f"  IDs   : {enc.ids}")
        print()


if __name__ == "__main__":
    main()