"""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
import argparse
import hashlib
import json
from pathlib import Path

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

ap = argparse.ArgumentParser()
ap.add_argument('--source', type=Path, default=Path.home() / 'run/cifar')
ap.add_argument('--out', type=Path, default=Path.home() / 'run/cifar/vehicles')
a = ap.parse_args()
src, out = a.source, a.out
out.mkdir(parents=True, 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)

# Native input format used by altai_infer; mirrors only augment the training split.
for split in ('train', 'val', 'test'):
    with np.load(out / f'{split}.npz') as saved:
        images, labels = saved['rgb'], saved['labels']
        images.tofile(out / f'{split}_images.bin')
        (out / f'{split}_labels.txt').write_text('\n'.join(map(str, labels.tolist())) + '\n')
        if split == 'train':
            images[:, :, :, ::-1].copy().tofile(out / 'train_flip_images.bin')
manifest = {'source': 'uoft-cs/cifar10 parquet', 'seed': 20261015, 'classes': NAMES,
            'split': 'first 4000 per class train, last 1000 validation, official test reserved',
            'files': {p.name: {'bytes': p.stat().st_size, 'sha256': hashlib.sha256(p.read_bytes()).hexdigest()}
                      for p in sorted(out.iterdir()) if p.suffix in {'.bin', '.txt', '.npz'}}}
(out / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')