"""Larger JPEG dataset with fresh validation/test splits for the second experiment."""
from concurrent.futures import ThreadPoolExecutor
import hashlib
from io import BytesIO
import json
from pathlib import Path
import time
import urllib.request

import numpy as np
from PIL import Image

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


def load(item):
    split, label, index = item
    path = cache / f'{split}-{label}-{index:04d}.jpg'
    if not path.exists():
        for attempt in range(5):
            try:
                with urllib.request.urlopen(f'{BASE}/{split}/{label}/{index:04d}.jpg', timeout=30) as response:
                    content = response.read()
                img = Image.open(BytesIO(content)).convert('RGB')
                assert img.size == (32, 32)
                path.write_bytes(content)
                break
            except Exception:
                if attempt == 4:
                    raise
                time.sleep(2 * (attempt + 1))
    return np.asarray(Image.open(path).convert('RGB')).transpose(2, 0, 1).reshape(3, 1024)


selected = {}
for split, original, lo, hi in [('train', 'train', 0, 2000), ('val', 'train', 4500, 5000),
                                ('test', 'test', 200, 700)]:
    items = [(original, label, i) for label in ('cat', 'dog') for i in range(lo, hi)]
    order = np.random.default_rng(20261015).permutation(len(items))
    selected[split] = [items[i] for i in order]
with ThreadPoolExecutor(max_workers=8) as pool:
    for split, items in selected.items():
        images = []
        for i, img in enumerate(pool.map(load, items)):
            images.append(img)
            if i % 500 == 0:
                print(split, i, '/', len(items), flush=True)
        labels = np.array([int(item[1] == 'dog') for item in items], dtype=np.uint8)
        np.savez_compressed(root / f'{split}-original.npz', rgb=np.stack(images), labels=labels)
        print(split, 'ready', flush=True)
assert not set(selected['train']) & set(selected['val'])
assert not any(index < 200 for _, _, index in selected['test'])
manifest = {'source': 'https://github.com/YoongiKim/CIFAR-10-images', 'commit': COMMIT,
            'variant': 'JPEG mirror, not canonical CIFAR bytes', 'seed': 20261015,
            'sizes': {s: len(items) for s, items in selected.items()}, 'source_files': selected,
            'fresh_validation': 'train indices 4500..4999 per class',
            'fresh_test': 'test indices 200..699 per class; excludes first experiment test 0..199'}
(root / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
print('V2 RGB dataset ready', flush=True)