"""Download CIFAR-10 JPEG mirror train images 2000–4499 for cat and dog (unused by v1/v2) into the same cache.

Same source and commit as cats-dogs/prepare_jpeg.py. Existing files are skipped.
usage: python3 fetch_more.py
"""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import time
import urllib.request

COMMIT = 'd7d05618061fe37d830fb0419063d113c575ccbc'
BASE = f'https://raw.githubusercontent.com/YoongiKim/CIFAR-10-images/{COMMIT}'
cache = Path('/home/int/run/cats-dogs/data/github-jpeg-cache')


def load(item):
    label, index = item
    path = cache / f'train-{label}-{index:04d}.jpg'
    if path.exists():
        return 0
    for attempt in range(6):
        try:
            with urllib.request.urlopen(f'{BASE}/train/{label}/{index:04d}.jpg', timeout=30) as r:
                data = r.read()
            tmp = path.with_suffix('.part')
            tmp.write_bytes(data)
            tmp.rename(path)
            return 1
        except Exception:
            if attempt == 5:
                raise
            time.sleep(3 * (attempt + 1))


items = [(label, i) for label in ('cat', 'dog') for i in range(2000, 4500)]
with ThreadPoolExecutor(max_workers=8) as pool:
    for n, _ in enumerate(pool.map(load, items)):
        if n % 500 == 0:
            print('done', n, flush=True)
print('all', len(items), flush=True)