"""Train BPE tokenizer on caption/VQA text data."""

import os
import json
from model.tokenizer import BPETokenizer


def collect_texts(flickr_dir: str = "data/flickr8k", vqa_dir: str = "data/vqav2") -> list[str]:
    """Collect all text from captions and VQA data."""
    texts = []

    # Flickr8k captions
    captions_file = os.path.join(flickr_dir, "captions.txt")
    if os.path.exists(captions_file):
        with open(captions_file) as f:
            for line in f:
                parts = line.strip().split("\t", 1)
                if len(parts) == 2:
                    texts.append(parts[1])

    # VQA questions + answers
    q_file = os.path.join(vqa_dir, "questions.json")
    a_file = os.path.join(vqa_dir, "annotations.json")
    if os.path.exists(q_file):
        with open(q_file) as f:
            for q in json.load(f)["questions"]:
                texts.append(q["question"])
    if os.path.exists(a_file):
        with open(a_file) as f:
            for a in json.load(f)["annotations"]:
                texts.append(a["multiple_choice_answer"])

    return texts


def main():
    texts = collect_texts()
    print(f"Collected {len(texts)} text samples")

    tokenizer = BPETokenizer(vocab_size=8192)
    print("Training BPE tokenizer...")
    tokenizer.train(texts)
    print(f"Vocabulary size: {len(tokenizer)}")

    # Save
    os.makedirs("checkpoints", exist_ok=True)
    tokenizer.save("checkpoints/tokenizer.json")
    print("Saved tokenizer to checkpoints/tokenizer.json")

    # Test
    test_texts = [
        "What color is the sky?",
        "a dog playing in the park",
        "How many people are there?",
    ]
    for text in test_texts:
        ids = tokenizer.encode(text)
        decoded = tokenizer.decode(ids)
        print(f"  '{text}' -> {ids[:10]}... -> '{decoded}'")


if __name__ == "__main__":
    main()