"""CIFAR-10 vehicles (airplane, automobile, ship, truck → 0..3) from the Hugging Face parquet copy (uoft-cs/cifar10).

train: first 4000 training images of each class (16000), val: the last 1000 of each class (4000),
test: the official test images of these classes (4000) — kept for one final evaluation.
usage: ~/anvenv/bin/python prepare_vehicles.py → /home/int/run/cifar/vehicles/{train,val,test}.npz
"""
import io
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
from PIL import Image

src, out = Path('/home/int/run/cifar'), Path('/home/int/run/cifar/vehicles')
out.mkdir(exist_ok=True)
CLASSES = {0: 0, 1: 1, 8: 2, 9: 3}  # CIFAR label → our label
NAMES = ['airplane', 'automobile', 'ship', 'truck']


def load(name):
    t = pq.read_table(src / f'hf_{name}.parquet').to_pydict()
    imgs, labels = [], []
    for img, lab in zip(t['img'], t['label']):
        if lab in CLASSES:
            imgs.append(np.asarray(Image.open(io.BytesIO(img['bytes'])).convert('RGB')).transpose(2, 0, 1))
            labels.append(CLASSES[lab])
    return np.stack(imgs), np.array(labels, np.uint8)


x, y = load('train')
idx = {c: np.flatnonzero(y == c) for c in range(4)}
train = np.concatenate([idx[c][:4000] for c in range(4)])
val = np.concatenate([idx[c][4000:] for c in range(4)])
rng = np.random.default_rng(20261015)
train, val = rng.permutation(train), rng.permutation(val)
np.savez_compressed(out / 'train.npz', rgb=x[train], labels=y[train])
np.savez_compressed(out / 'val.npz', rgb=x[val], labels=y[val])
xt, yt = load('test')
p = rng.permutation(len(yt))
np.savez_compressed(out / 'test.npz', rgb=xt[p], labels=yt[p])
print('train', len(train), np.bincount(y[train]), 'val', len(val), 'test', len(yt), np.bincount(yt), NAMES)