scripts/download_all_data.py

#!/usr/bin/env python3
"""
Download ALL training data with REAL images for ArcisVLM.

This is the SINGLE entry point for data preparation. Run this before ANY training.
Every dataset downloads actual photographs — no dummies, no noise, no fallbacks.

Usage:
    python3 scripts/download_all_data.py                    # Download all stages
    python3 scripts/download_all_data.py --stage 1          # Stage 1 only
    python3 scripts/download_all_data.py --stage 2          # Stage 2 only
    python3 scripts/download_all_data.py --stage 3          # Stage 3 only
    python3 scripts/download_all_data.py --max-per-dataset 10000  # Cap per dataset (for testing)

Disk requirements:
    Stage 1: ~50GB  (500K image-caption pairs)
    Stage 2: ~150GB (2M VQA samples with images)
    Stage 3: ~20GB  (200K detection/surveillance)
    Total:   ~220GB
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path


def _save_image(image, path):
    """Save a PIL image to disk. Returns True on success."""
    try:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        if hasattr(image, 'save'):
            image.convert("RGB").save(path, quality=85)
            return True
    except Exception:
        return False
    return False


def _download_hf_dataset(name, hf_repo, split, output_dir, max_samples,
                          q_key=None, a_key=None, img_key="image",
                          config_name=None, caption_mode=False):
    """Download a HuggingFace dataset with real images to JSONL."""
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")
    img_dir = os.path.join(output_dir, "images", name)

    # Resume support: check existing
    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    print(f"  [{name}] Downloading from {hf_repo} ({split}, max {max_samples})...")
    os.makedirs(img_dir, exist_ok=True)

    from datasets import load_dataset

    kwargs = {"streaming": True}
    if config_name:
        ds = load_dataset(hf_repo, config_name, split=split, **kwargs)
    else:
        ds = load_dataset(hf_repo, split=split, **kwargs)

    count = 0
    skipped_no_image = 0

    with open(jsonl_path, "w") as f:
        for item in ds:
            if count >= max_samples:
                break

            # Extract text
            if caption_mode:
                # Image captioning format
                question = "Describe this image."
                answer = item.get("caption", item.get("text", item.get("sentence", "")))
                if isinstance(answer, list):
                    answer = answer[0] if answer else ""
            elif q_key and a_key:
                question = item.get(q_key, "")
                answer = item.get(a_key, "")
                if isinstance(answer, list):
                    if answer and isinstance(answer[0], dict):
                        answer = answer[0].get("answer", str(answer[0]))
                    elif answer:
                        answer = str(answer[0])
                    else:
                        answer = ""
            else:
                continue

            if not answer:
                continue

            # Get and save image — REQUIRED, skip if no image
            image = item.get(img_key)
            if image is None:
                skipped_no_image += 1
                continue

            img_filename = f"{name}_{count:07d}.jpg"
            img_path = os.path.join(img_dir, img_filename)

            if not _save_image(image, img_path):
                skipped_no_image += 1
                continue

            sample = {
                "question": str(question).strip(),
                "answer": str(answer).strip(),
                "image_path": img_path,
                "dataset": name,
            }
            f.write(json.dumps(sample) + "\n")
            count += 1

            if count % 5000 == 0:
                print(f"    [{name}] {count:,} samples (skipped {skipped_no_image} without images)...")

    print(f"  [{name}] Done: {count:,} samples saved (skipped {skipped_no_image})")
    return count


def _download_coco_detection(output_dir, max_samples):
    """Download COCO detection dataset with real images and human-readable category names."""
    # COCO category ID → human-readable name mapping
    COCO_CATEGORIES = {
        0: "person", 1: "bicycle", 2: "car", 3: "motorcycle", 4: "airplane",
        5: "bus", 6: "train", 7: "truck", 8: "boat", 9: "traffic light",
        10: "fire hydrant", 11: "stop sign", 12: "parking meter", 13: "bench",
        14: "bird", 15: "cat", 16: "dog", 17: "horse", 18: "sheep",
        19: "cow", 20: "elephant", 21: "bear", 22: "zebra", 23: "giraffe",
        24: "backpack", 25: "umbrella", 26: "handbag", 27: "tie", 28: "suitcase",
        29: "frisbee", 30: "skis", 31: "snowboard", 32: "sports ball", 33: "kite",
        34: "baseball bat", 35: "baseball glove", 36: "skateboard", 37: "surfboard",
        38: "tennis racket", 39: "bottle", 40: "wine glass", 41: "cup",
        42: "fork", 43: "knife", 44: "spoon", 45: "bowl", 46: "banana",
        47: "apple", 48: "sandwich", 49: "orange", 50: "broccoli",
        51: "carrot", 52: "hot dog", 53: "pizza", 54: "donut", 55: "cake",
        56: "chair", 57: "couch", 58: "potted plant", 59: "bed",
        60: "dining table", 61: "toilet", 62: "tv", 63: "laptop",
        64: "mouse", 65: "remote", 66: "keyboard", 67: "cell phone",
        68: "microwave", 69: "oven", 70: "toaster", 71: "sink",
        72: "refrigerator", 73: "book", 74: "clock", 75: "vase",
        76: "scissors", 77: "teddy bear", 78: "hair drier", 79: "toothbrush",
    }

    name = "coco_detect"
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")
    img_dir = os.path.join(output_dir, "images", name)
    os.makedirs(img_dir, exist_ok=True)

    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    print(f"  [{name}] Downloading from detection-datasets/coco (max {max_samples})...")
    from datasets import load_dataset

    ds = load_dataset("detection-datasets/coco", split="train", streaming=True)
    count = 0

    with open(jsonl_path, "w") as f:
        for item in ds:
            if count >= max_samples:
                break

            image = item.get("image")
            if image is None:
                continue

            # Extract object categories as HUMAN-READABLE NAMES (not integer IDs!)
            objects = item.get("objects", {})
            categories = objects.get("category", []) if isinstance(objects, dict) else []
            if isinstance(categories, list) and categories:
                # Map IDs to names, deduplicate, preserve order
                names = []
                seen = set()
                for c in categories[:15]:
                    name_str = COCO_CATEGORIES.get(int(c), f"object")
                    if name_str not in seen:
                        names.append(name_str)
                        seen.add(name_str)
                caption = ", ".join(names)
            else:
                caption = "a photograph"

            img_filename = f"coco_detect_{count:07d}.jpg"
            img_path = os.path.join(img_dir, img_filename)

            if not _save_image(image, img_path):
                continue

            # Diverse questions for richer training
            import random
            q_templates = [
                f"What objects are in this image?",
                f"Describe what you see.",
                f"List the objects visible in this scene.",
                f"What is in this photograph?",
                f"How many objects can you identify?",
            ]
            question = q_templates[count % len(q_templates)]

            # Rich answer format
            if len(names) == 1:
                answer = f"I can see a {names[0]}."
            elif len(names) == 2:
                answer = f"I can see a {names[0]} and a {names[1]}."
            else:
                answer = f"I can see {', '.join(names[:-1])}, and {names[-1]}."

            sample = {
                "question": question,
                "answer": answer,
                "image_path": img_path,
                "dataset": "coco_detect",
            }
            f.write(json.dumps(sample) + "\n")
            count += 1

            if count % 5000 == 0:
                print(f"    [coco_detect] {count:,} samples...")

    print(f"  [{name}] Done: {count:,} samples")
    return count


def _download_scienceqa(output_dir, max_samples):
    """Download ScienceQA with real images."""
    name = "scienceqa"
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")
    img_dir = os.path.join(output_dir, "images", name)
    os.makedirs(img_dir, exist_ok=True)

    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    print(f"  [{name}] Downloading from derek-thomas/ScienceQA (max {max_samples})...")
    from datasets import load_dataset

    ds = load_dataset("derek-thomas/ScienceQA", split="train", streaming=True)
    count = 0
    skipped = 0

    with open(jsonl_path, "w") as f:
        for item in ds:
            if count >= max_samples:
                break

            image = item.get("image")
            if image is None:
                skipped += 1
                continue

            question = item.get("question", "")
            answer_idx = item.get("answer", 0)
            choices = item.get("choices", [])

            # Map answer index to text
            if isinstance(answer_idx, int) and choices and answer_idx < len(choices):
                answer = str(choices[answer_idx])
            else:
                answer = str(answer_idx)

            if not question or not answer:
                continue

            img_filename = f"{name}_{count:07d}.jpg"
            img_path = os.path.join(img_dir, img_filename)

            if not _save_image(image, img_path):
                continue

            sample = {
                "question": question,
                "answer": answer,
                "image_path": img_path,
                "dataset": name,
            }
            f.write(json.dumps(sample) + "\n")
            count += 1

            if count % 2000 == 0:
                print(f"    [{name}] {count:,} samples (skipped {skipped} without images)...")

    print(f"  [{name}] Done: {count:,} samples (skipped {skipped})")
    return count


def download_stage1(output_dir, max_per_dataset):
    """Download Stage 1 JEPA pretraining data (image-caption pairs with real photos)."""
    print("\n" + "=" * 60)
    print("STAGE 1: JEPA Pretraining Data (image-caption pairs)")
    print("=" * 60)

    os.makedirs(output_dir, exist_ok=True)
    total = 0

    datasets = [
        # COCO detection — 120K real photos (extract object list as caption)
        {"name": "coco_detect", "hf_repo": "detection-datasets/coco", "split": "train",
         "target": 120000, "custom_loader": "coco"},
        # ScienceQA — 12K real images with questions
        {"name": "scienceqa", "hf_repo": "derek-thomas/ScienceQA", "split": "train",
         "target": 12000, "custom_loader": "scienceqa"},
    ]

    for ds in datasets:
        target = min(ds["target"], max_per_dataset) if max_per_dataset else ds["target"]
        custom = ds.get("custom_loader")
        try:
            if custom == "coco":
                count = _download_coco_detection(output_dir, target)
            elif custom == "scienceqa":
                count = _download_scienceqa(output_dir, target)
            else:
                count = _download_hf_dataset(
                    name=ds["name"], hf_repo=ds["hf_repo"], split=ds["split"],
                    output_dir=output_dir, max_samples=target,
                    img_key=ds.get("img_key", "image"),
                    config_name=ds.get("config_name"),
                    caption_mode=ds.get("caption_mode", False),
                )
            total += count
        except Exception as e:
            print(f"  [{ds['name']}] FAILED: {e}")

    print(f"\nStage 1 total: {total:,} samples")
    return total


def _download_coco_captions(output_dir, max_samples):
    """Download COCO images with FULL SENTENCE captions (for Caption agent)."""
    name = "coco_captions"
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")
    img_dir = os.path.join(output_dir, "images", name)
    os.makedirs(img_dir, exist_ok=True)

    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    # Same COCO dataset but generate descriptive captions instead of object lists
    COCO_CATS = {
        0: "person", 1: "bicycle", 2: "car", 3: "motorcycle", 4: "airplane",
        5: "bus", 6: "train", 7: "truck", 8: "boat", 9: "traffic light",
        10: "fire hydrant", 11: "stop sign", 12: "parking meter", 13: "bench",
        14: "bird", 15: "cat", 16: "dog", 17: "horse", 18: "sheep",
        19: "cow", 20: "elephant", 21: "bear", 22: "zebra", 23: "giraffe",
        24: "backpack", 25: "umbrella", 26: "handbag", 27: "tie", 28: "suitcase",
        29: "frisbee", 30: "skis", 31: "snowboard", 32: "sports ball", 33: "kite",
        34: "baseball bat", 35: "baseball glove", 36: "skateboard", 37: "surfboard",
        38: "tennis racket", 39: "bottle", 40: "wine glass", 41: "cup",
        42: "fork", 43: "knife", 44: "spoon", 45: "bowl", 46: "banana",
        47: "apple", 48: "sandwich", 49: "orange", 50: "broccoli",
        51: "carrot", 52: "hot dog", 53: "pizza", 54: "donut", 55: "cake",
        56: "chair", 57: "couch", 58: "potted plant", 59: "bed",
        60: "dining table", 61: "toilet", 62: "tv", 63: "laptop",
        64: "mouse", 65: "remote", 66: "keyboard", 67: "cell phone",
        68: "microwave", 69: "oven", 70: "toaster", 71: "sink",
        72: "refrigerator", 73: "book", 74: "clock", 75: "vase",
        76: "scissors", 77: "teddy bear", 78: "hair drier", 79: "toothbrush",
    }
    import random

    print(f"  [{name}] Generating descriptive captions from COCO (max {max_samples})...")
    from datasets import load_dataset
    ds = load_dataset("detection-datasets/coco", split="train", streaming=True)
    count = 0

    templates = [
        "This image shows {scene}.",
        "A scene containing {scene}.",
        "In this photograph, there are {scene}.",
        "The image depicts {scene}.",
        "Visible in this scene: {scene}.",
    ]

    with open(jsonl_path, "w") as f:
        for item in ds:
            if count >= max_samples:
                break
            image = item.get("image")
            if image is None:
                continue
            objects = item.get("objects", {})
            categories = objects.get("category", []) if isinstance(objects, dict) else []
            if not categories:
                continue

            names = list(dict.fromkeys(COCO_CATS.get(int(c), "object") for c in categories[:10]))
            # Count objects
            from collections import Counter
            cat_counts = Counter(COCO_CATS.get(int(c), "object") for c in categories)
            parts = []
            for obj, cnt in cat_counts.most_common(5):
                if cnt == 1:
                    parts.append(f"a {obj}")
                else:
                    parts.append(f"{cnt} {obj}s")
            if len(parts) > 1:
                scene = ", ".join(parts[:-1]) + f", and {parts[-1]}"
            else:
                scene = parts[0]

            caption = random.choice(templates).format(scene=scene)

            img_filename = f"{name}_{count:07d}.jpg"
            img_path = os.path.join(img_dir, img_filename)
            if not _save_image(image, img_path):
                continue

            f.write(json.dumps({
                "question": "Describe this image in detail.",
                "answer": caption,
                "image_path": img_path,
                "dataset": name,
            }) + "\n")
            count += 1
            if count % 5000 == 0:
                print(f"    [{name}] {count:,}...")

    print(f"  [{name}] Done: {count:,}")
    return count


def _download_coco_count(output_dir, max_samples):
    """Derive counting data from COCO detection (for Count agent)."""
    name = "coco_count"
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")
    img_dir = os.path.join(output_dir, "images", name)
    os.makedirs(img_dir, exist_ok=True)

    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    COCO_CATS = {
        0: "person", 1: "bicycle", 2: "car", 3: "motorcycle", 4: "airplane",
        5: "bus", 6: "train", 7: "truck", 8: "boat", 9: "traffic light",
        14: "bird", 15: "cat", 16: "dog", 17: "horse", 18: "sheep",
        19: "cow", 20: "elephant", 22: "zebra", 23: "giraffe",
        56: "chair", 57: "couch", 59: "bed", 60: "dining table",
    }
    import random

    print(f"  [{name}] Generating counting data from COCO (max {max_samples})...")
    from datasets import load_dataset
    ds = load_dataset("detection-datasets/coco", split="train", streaming=True)
    count = 0

    q_templates = [
        "How many {obj} are in this image?",
        "Count the {obj}.",
        "How many {obj} can you see?",
    ]

    with open(jsonl_path, "w") as f:
        for item in ds:
            if count >= max_samples:
                break
            image = item.get("image")
            if image is None:
                continue
            objects = item.get("objects", {})
            categories = objects.get("category", []) if isinstance(objects, dict) else []
            if not categories:
                continue

            from collections import Counter
            cat_counts = Counter(COCO_CATS.get(int(c)) for c in categories if int(c) in COCO_CATS)
            if not cat_counts:
                continue

            img_filename = f"{name}_{count:07d}.jpg"
            img_path = os.path.join(img_dir, img_filename)
            if not _save_image(image, img_path):
                continue

            # Generate counting Q&A for the most common object
            obj, cnt = cat_counts.most_common(1)[0]
            question = random.choice(q_templates).format(obj=obj)

            f.write(json.dumps({
                "question": question,
                "answer": str(cnt),
                "image_path": img_path,
                "dataset": name,
            }) + "\n")
            count += 1
            if count % 5000 == 0:
                print(f"    [{name}] {count:,}...")

    print(f"  [{name}] Done: {count:,}")
    return count


def _generate_surveillance_qa(output_dir, max_samples):
    """Generate surveillance-domain Q&A for Alert + Reason agents."""
    name = "surveillance_qa"
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")

    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    import random
    print(f"  [{name}] Generating surveillance Q&A (max {max_samples})...")

    locations = ["parking lot", "lobby", "hallway", "entrance", "warehouse", "street",
                 "intersection", "loading dock", "stairwell", "elevator area", "gate", "rooftop"]
    objects = ["person", "car", "truck", "bicycle", "motorcycle", "bus", "dog",
               "backpack", "suitcase", "umbrella", "skateboard"]
    actions = ["walking", "running", "standing", "sitting", "carrying a bag",
               "talking on phone", "entering", "exiting", "loitering", "crossing the road"]
    times = ["daytime", "nighttime", "dawn", "dusk", "overcast conditions"]
    severities = ["LOW", "MEDIUM", "HIGH", "CRITICAL"]

    alert_scenarios = [
        ("Person detected in restricted area near {loc}.", "HIGH", "Unauthorized access"),
        ("Unattended bag found near {loc}.", "HIGH", "Suspicious object"),
        ("Person loitering at {loc} for extended period.", "MEDIUM", "Loitering"),
        ("Vehicle parked in no-parking zone at {loc}.", "LOW", "Parking violation"),
        ("Person running in {loc}.", "MEDIUM", "Unusual behavior"),
        ("Group gathering detected at {loc}.", "LOW", "Crowd formation"),
        ("Person climbing fence near {loc}.", "HIGH", "Trespassing attempt"),
        ("Fight detected between two people at {loc}.", "CRITICAL", "Violence"),
        ("Person fallen on ground at {loc}.", "CRITICAL", "Fall detection"),
        ("Smoke detected near {loc}.", "CRITICAL", "Fire hazard"),
        ("No anomalies detected. Normal activity at {loc}.", "LOW", "All clear"),
        ("Tailgating detected at {loc} entrance.", "MEDIUM", "Access control violation"),
    ]

    count = 0
    with open(jsonl_path, "w") as f:
        for _ in range(max_samples):
            loc = random.choice(locations)
            scenario = random.choice(alert_scenarios)
            description = scenario[0].format(loc=loc)
            severity = scenario[1]
            category = scenario[2]

            # Alert format
            if random.random() < 0.4:
                question = "Is there any security concern in this scene?"
                answer = f"ALERT: {description} Severity: {severity}. Category: {category}."
            # Reason format
            elif random.random() < 0.7:
                n_people = random.randint(0, 8)
                n_vehicles = random.randint(0, 5)
                time = random.choice(times)
                question = "Analyze this scene and describe what you observe."
                answer = (f"Scene analysis: The {loc} during {time}. "
                         f"I observe {n_people} people and {n_vehicles} vehicles. "
                         f"{description}")
            # Count format
            else:
                n = random.randint(0, 15)
                obj = random.choice(objects)
                question = f"How many {obj}s are in this scene?"
                answer = str(n)

            f.write(json.dumps({
                "question": question,
                "answer": answer,
                "dataset": name,
            }) + "\n")
            count += 1

    print(f"  [{name}] Done: {count:,} (text-only, no images)")
    return count


def _download_ocrvqa(output_dir, max_samples):
    """Download OCR-VQA with proper question-answer extraction.

    OCR-VQA has 'questions' (list) and 'answers' (list), not singular fields.
    Each image has ~5 Q&A pairs about book covers.
    """
    name = "ocrvqa"
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")
    img_dir = os.path.join(output_dir, "images", name)
    os.makedirs(img_dir, exist_ok=True)

    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    print(f"  [{name}] Downloading from howard-hou/OCR-VQA (max {max_samples})...")
    from datasets import load_dataset
    ds = load_dataset("howard-hou/OCR-VQA", split="train", streaming=True)
    count = 0

    with open(jsonl_path, "w") as f:
        for item in ds:
            if count >= max_samples:
                break

            image = item.get("image")
            if image is None:
                continue

            # Extract Q&A pairs from LISTS
            questions = item.get("questions", [])
            answers = item.get("answers", [])

            if not questions or not answers:
                continue

            img_filename = f"{name}_{count:07d}.jpg"
            img_path = os.path.join(img_dir, img_filename)
            if not _save_image(image, img_path):
                continue

            # Write each Q&A pair as a separate sample
            for q, a in zip(questions, answers):
                if q and a and len(q) > 3 and len(a) > 0:
                    f.write(json.dumps({
                        "question": str(q),
                        "answer": str(a),
                        "image_path": img_path,
                        "dataset": name,
                    }) + "\n")
                    count += 1
                    if count >= max_samples:
                        break

            if count % 5000 == 0:
                print(f"    [{name}] {count:,}...")

    print(f"  [{name}] Done: {count:,}")
    return count


def _download_aokvqa(output_dir, max_samples):
    """Download A-OKVQA with proper answer extraction.

    A-OKVQA 'direct_answers' is a LIST of 10 annotator answers.
    We take the most common answer (majority vote).
    """
    name = "aokvqa"
    jsonl_path = os.path.join(output_dir, f"{name}.jsonl")
    img_dir = os.path.join(output_dir, "images", name)
    os.makedirs(img_dir, exist_ok=True)

    if os.path.exists(jsonl_path):
        with open(jsonl_path) as f:
            existing = sum(1 for _ in f)
        if existing >= max_samples * 0.9:
            print(f"  [{name}] Already have {existing} samples, skipping")
            return existing

    print(f"  [{name}] Downloading from HuggingFaceM4/A-OKVQA (max {max_samples})...")
    from datasets import load_dataset
    from collections import Counter
    ds = load_dataset("HuggingFaceM4/A-OKVQA", split="train", streaming=True)
    count = 0

    with open(jsonl_path, "w") as f:
        for item in ds:
            if count >= max_samples:
                break

            image = item.get("image")
            if image is None:
                continue

            question = item.get("question", "")
            direct_answers = item.get("direct_answers", [])

            if not question or not direct_answers:
                continue

            # Take majority vote answer (most common in the list)
            if isinstance(direct_answers, list):
                answer_counts = Counter(str(a) for a in direct_answers if a)
                if answer_counts:
                    answer = answer_counts.most_common(1)[0][0]
                else:
                    continue
            else:
                answer = str(direct_answers)

            if not answer or len(answer) < 1:
                continue

            img_filename = f"{name}_{count:07d}.jpg"
            img_path = os.path.join(img_dir, img_filename)
            if not _save_image(image, img_path):
                continue

            f.write(json.dumps({
                "question": str(question),
                "answer": answer,
                "image_path": img_path,
                "dataset": name,
            }) + "\n")
            count += 1

            if count % 5000 == 0:
                print(f"    [{name}] {count:,}...")

    print(f"  [{name}] Done: {count:,}")
    return count


def download_stage2(output_dir, max_per_dataset):
    """Download Stage 2 instruction tuning data (VQA with real images)."""
    print("\n" + "=" * 60)
    print("STAGE 2: Instruction Tuning Data (VQA with real images)")
    print("=" * 60)

    os.makedirs(output_dir, exist_ok=True)
    total = 0

    datasets = [
        # VQA agent: short factual answers
        {"name": "vqav2", "hf_repo": "merve/vqav2-small", "split": "validation",
         "q_key": "question", "a_key": "multiple_choice_answer", "target": 40000,
         "img_key": "image"},
        # VQA + Reason: scene graph spatial reasoning
        {"name": "gqa", "hf_repo": "lmms-lab/GQA", "split": "train",
         "q_key": "question", "a_key": "answer", "target": 200000},
        # OCR agent: text in natural images
        {"name": "textvqa", "hf_repo": "lmms-lab/textvqa", "split": "train",
         "q_key": "question", "a_key": "answers", "target": 34000},
        # OCR agent: document/book text (fields are "questions" and "answers" LISTS)
        {"name": "ocrvqa", "hf_repo": "howard-hou/OCR-VQA", "split": "train",
         "target": 207000, "custom_loader": "ocrvqa"},
        # Reason agent: knowledge-requiring questions (direct_answers is a LIST)
        {"name": "aokvqa", "hf_repo": "HuggingFaceM4/A-OKVQA", "split": "train",
         "target": 17000, "custom_loader": "aokvqa"},
        # VQA + Reason: science questions with images
        {"name": "scienceqa", "hf_repo": "derek-thomas/ScienceQA", "split": "train",
         "target": 12000, "custom_loader": "scienceqa"},
        # Caption agent: COCO detection images with FULL SENTENCE captions
        {"name": "coco_captions", "hf_repo": "detection-datasets/coco", "split": "train",
         "target": 50000, "custom_loader": "coco_captions"},
        # Count agent: derived counting data from COCO
        {"name": "coco_count", "hf_repo": "detection-datasets/coco", "split": "train",
         "target": 50000, "custom_loader": "coco_count"},
        # NOTE: surveillance_qa is TEXT-ONLY (no images) — moved to Stage 3
        # Do NOT include text-only data in Stage 2 (zero-tolerance: crashes on missing images)
    ]

    for ds in datasets:
        target = min(ds["target"], max_per_dataset) if max_per_dataset else ds["target"]
        custom = ds.get("custom_loader")
        try:
            if custom == "scienceqa":
                count = _download_scienceqa(output_dir, target)
            elif custom == "coco":
                count = _download_coco_detection(output_dir, target)
            elif custom == "coco_captions":
                count = _download_coco_captions(output_dir, target)
            elif custom == "coco_count":
                count = _download_coco_count(output_dir, target)
            elif custom == "surveillance_qa":
                count = _generate_surveillance_qa(output_dir, target)
            elif custom == "ocrvqa":
                count = _download_ocrvqa(output_dir, target)
            elif custom == "aokvqa":
                count = _download_aokvqa(output_dir, target)
            else:
                count = _download_hf_dataset(
                    name=ds["name"], hf_repo=ds["hf_repo"], split=ds["split"],
                    output_dir=output_dir, max_samples=target,
                    q_key=ds.get("q_key"), a_key=ds.get("a_key"),
                    img_key=ds.get("img_key", "image"),
                    config_name=ds.get("config_name"),
                )
            total += count
        except Exception as e:
            print(f"  [{ds['name']}] FAILED: {e}")

    print(f"\nStage 2 total: {total:,} samples")
    return total


def download_stage3(output_dir, max_per_dataset):
    """Download Stage 3 domain fine-tuning data."""
    print("\n" + "=" * 60)
    print("STAGE 3: Domain Fine-Tuning Data (detection/surveillance)")
    print("=" * 60)

    # Use existing download_stage3_data.py
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    try:
        from download_stage3_data import main as download_stage3_main
        # Temporarily set output dir
        import download_stage3_data
        old_main = download_stage3_data.main
        download_stage3_data.main = lambda: None  # Prevent double execution

        from download_stage3_data import (
            download_coco_detection, download_visdrone,
            download_activitynet_captions, download_ucf_crime,
            download_surveillance_vqa
        )

        os.makedirs(output_dir, exist_ok=True)
        total = 0
        cap = max_per_dataset or 200000

        total += download_coco_detection(output_dir, min(118000, cap))
        total += download_visdrone(output_dir, min(10000, cap))
        total += download_activitynet_captions(output_dir, min(100000, cap))
        total += download_ucf_crime(output_dir, min(1900, cap))
        total += download_surveillance_vqa(output_dir, min(50000, cap))

        print(f"\nStage 3 total: {total:,} samples")
        return total
    except ImportError:
        print("  [WARN] download_stage3_data.py not importable, running as subprocess")
        import subprocess
        subprocess.run([sys.executable, "scripts/download_stage3_data.py"], check=True)
        return 0


def validate_data(data_dir, stage_name, min_samples=100):
    """Validate downloaded data has real images."""
    jsonl_files = list(Path(data_dir).glob("*.jsonl"))
    if not jsonl_files:
        raise RuntimeError(f"No JSONL files found in {data_dir}")

    # Count total samples across all JSONL files
    total = 0
    for jf in jsonl_files:
        with open(jf) as f:
            total += sum(1 for _ in f)

    # Spot-check that images exist for a sample of entries
    images_checked = 0
    images_found = 0
    for jf in jsonl_files:
        with open(jf) as f:
            for line in f:
                item = json.loads(line)
                if "image_path" in item:
                    images_checked += 1
                    if os.path.exists(item["image_path"]):
                        images_found += 1
                if images_checked >= 50:
                    break
        if images_checked >= 50:
            break

    print(f"\n  [{stage_name}] Validation: {total} total samples, {images_found}/{images_checked} spot-checked have real images")
    if total < min_samples:
        raise RuntimeError(f"FATAL: Only {total} samples in {data_dir}. Need at least {min_samples}.")
    if images_found == 0:
        raise RuntimeError(f"FATAL: No real images found in spot check of {data_dir}.")

    print(f"  [{stage_name}] PASSED")


def main():
    parser = argparse.ArgumentParser(description="Download ALL ArcisVLM Training Data")
    parser.add_argument("--stage", type=int, choices=[1, 2, 3], default=None,
                        help="Download specific stage only (default: all)")
    parser.add_argument("--max-per-dataset", type=int, default=None,
                        help="Cap samples per dataset (for testing)")
    parser.add_argument("--skip-validation", action="store_true")
    args = parser.parse_args()

    start = time.time()
    print("=" * 60)
    print("ArcisVLM — Download ALL Training Data (REAL images only)")
    print("=" * 60)

    stages = [args.stage] if args.stage else [1, 2, 3]

    if 1 in stages:
        download_stage1("data/downloads/stage1", args.max_per_dataset)
        if not args.skip_validation:
            validate_data("data/downloads/stage1", "Stage 1")

    if 2 in stages:
        download_stage2("data/downloads/stage2_fullscale", args.max_per_dataset)
        if not args.skip_validation:
            validate_data("data/downloads/stage2_fullscale", "Stage 2")

    if 3 in stages:
        download_stage3("data/downloads/stage3", args.max_per_dataset)
        if not args.skip_validation:
            validate_data("data/downloads/stage3", "Stage 3")

    elapsed = time.time() - start
    print(f"\n{'=' * 60}")
    print(f"Download complete in {elapsed/60:.0f} minutes")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()