"""
quick_test_flickr8k.py

Quick script to:
1. Extract your Flickr8k zip (if not already extracted)
2. Inspect the dataset structure
3. Run a pretrained BLIP model on a few sample images
   (No training required — instant results!)

Usage:
    python quick_test_flickr8k.py
"""

import os
import zipfile
import random
from pathlib import Path
from PIL import Image
from tqdm import tqdm
import torch
from transformers import BlipProcessor, BlipForConditionalGeneration


# ============================
# CONFIG — adjust if needed
# ============================
ZIP_PATH = "archive.zip"
EXTRACT_TO = "./flickr8k_extracted"
NUM_SAMPLES = 5  # How many random images to caption
MODEL_NAME = "Salesforce/blip-image-captioning-base"


def extract_zip(zip_path, extract_to):
    """Extract zip if not already extracted."""
    extract_path = Path(extract_to)
    if extract_path.exists() and any(extract_path.iterdir()):
        print(f"[INFO] Already extracted at: {extract_path}")
        return str(extract_path)

    print(f"[INFO] Extracting {zip_path} → {extract_path} ...")
    extract_path.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path, 'r') as zf:
        zf.extractall(extract_path)
    print("[INFO] Extraction complete.")
    return str(extract_path)


def find_images_folder(root):
    """Find the folder containing .jpg images."""
    root = Path(root)

    # Common names
    for name in ["Flicker8k_Dataset", "Flickr8k_Dataset", "images", "Images"]:
        p = root / name
        if p.exists() and p.is_dir():
            return p
        # Also check one level deep
        for sub in root.iterdir():
            if sub.is_dir():
                p2 = sub / name
                if p2.exists() and p2.is_dir():
                    return p2

    # Fallback: find any folder with .jpg files
    for sub in root.rglob("*.jpg"):
        return sub.parent

    return None


def main():
    print("=" * 60)
    print("  FLICKR8K QUICK TEST — Pretrained BLIP")
    print("=" * 60)

    # 1. Extract
    extracted = extract_zip(ZIP_PATH, EXTRACT_TO)
    print(f"[INFO] Extracted root: {extracted}")

    # 2. Find images
    img_folder = find_images_folder(extracted)
    if img_folder is None:
        print("[ERROR] Could not find image folder. Check your zip structure.")
        return

    print(f"[INFO] Image folder: {img_folder}")

    # 3. Collect images
    all_images = sorted([f for f in img_folder.iterdir() if f.suffix.lower() in ('.jpg', '.jpeg', '.png')])
    print(f"[INFO] Found {len(all_images)} images")

    if len(all_images) == 0:
        print("[ERROR] No images found.")
        return

    # Sample random images
    samples = random.sample(all_images, min(NUM_SAMPLES, len(all_images)))
    print(f"[INFO] Testing on {len(samples)} random images...\n")

    # 4. Load model
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[INFO] Loading BLIP model on {device} ...")
    processor = BlipProcessor.from_pretrained(MODEL_NAME)
    model = BlipForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
    print("[INFO] Model loaded. Generating captions...\n")

    # 5. Generate captions
    for img_path in samples:
        image = Image.open(img_path).convert("RGB")
        inputs = processor(images=image, return_tensors="pt").to(device)

        with torch.no_grad():
            output = model.generate(**inputs, max_length=50, num_beams=5)

        caption = processor.decode(output[0], skip_special_tokens=True)
        print(f"  {img_path.name}")
        print(f"    → \"{caption}\"\n")

    print("[DONE] Test complete!")


if __name__ == "__main__":
    main()