"""Cats/dogs v3 data: train = CIFAR-10 JPEG mirror train 0000–4499 per class (9000 images),
validation = train 4500–4999 per class (same images as v2 validation). Test is not touched.

usage: ~/anvenv/bin/python prepare_big.py   → /home/int/run/cats-dogs/data-v3/{train,val}.npz
"""
from pathlib import Path

import numpy as np
from PIL import Image

cache = Path('/home/int/run/cats-dogs/data/github-jpeg-cache')
out = Path('/home/int/run/cats-dogs/data-v3')
out.mkdir(exist_ok=True)


def load(lo, hi, seed):
    items = [(label, cls, i) for cls, label in enumerate(('cat', 'dog')) for i in range(lo, hi)]
    order = np.random.default_rng(seed).permutation(len(items))
    items = [items[j] for j in order]
    rgb = np.stack([np.asarray(Image.open(cache / f'train-{label}-{i:04d}.jpg').convert('RGB')).transpose(2, 0, 1)
                    for label, _, i in items])
    return rgb, np.array([c for _, c, _ in items], dtype=np.uint8), [f'train/{l}/{i:04d}' for l, _, i in items]


for name, lo, hi in [('train', 0, 4500), ('val', 4500, 5000)]:
    rgb, labels, names = load(lo, hi, 20261015)
    np.savez_compressed(out / f'{name}.npz', rgb=rgb, labels=labels, names=np.array(names))
    print(name, rgb.shape, np.bincount(labels))