"""Reproducible pilot from a commit-pinned JPEG mirror of CIFAR-10."""
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
import hashlib
import json
from pathlib import Path
import time
import urllib.request

import numpy as np
from PIL import Image
from prepare_data import encode

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


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


selected = {}
for split, original, lo, hi in [('train', 'train', 0, 400), ('val', 'train', 400, 500),
                                ('test', 'test', 0, 200)]:
    items = [(original, label, i) for label in ('cat', 'dog') for i in range(lo, hi)]
    order = np.random.default_rng(20261014).permutation(len(items))
    selected[split] = [items[i] for i in order]
arrays, labels = {}, {}
with ThreadPoolExecutor(max_workers=8) as pool:
    for split, items in selected.items():
        arrays[split] = np.stack(list(pool.map(load, items)))
        labels[split] = np.array([int(item[1] == 'dog') for item in items], dtype=np.uint8)
        encode(arrays[split]).tofile(root / f'{split}.bin')
        np.savetxt(root / f'{split}.target', labels[split], fmt='%d')
        np.savez_compressed(root / f'{split}-original.npz', rgb=arrays[split], labels=labels[split])
        print(split, len(items), 'ready', flush=True)
for name, first, second in [('smoke', ('train', 100), ('val', 100)),
                             ('validation', ('train', 800), ('val', 200)),
                             ('train_test', ('train', 800), ('test', 400))]:
    encode(np.concatenate([arrays[first[0]][:first[1]], arrays[second[0]][:second[1]]])).tofile(root / f'{name}.bin')
    np.savetxt(root / f'{name}.target', np.concatenate([labels[first[0]][:first[1]], labels[second[0]][:second[1]]]), fmt='%d')
assert not set(selected['train']) & set(selected['val'])
manifest = {'source': 'https://github.com/YoongiKim/CIFAR-10-images', 'commit': COMMIT,
            'variant': 'JPEG mirror of CIFAR-10; not original binary archive',
            'labels': {'0': 'cat', '1': 'dog'}, 'seed': 20261014,
            'train': 800, 'validation': 200, 'test': 400, 'features': 2048,
            'encoding': 'integer grayscale BT.601, concatenate brightness and darkness',
            'source_files': selected, 'train_validation_disjoint': True,
            'encoded_sha256': {s: hashlib.sha256((root / f'{s}.bin').read_bytes()).hexdigest() for s in selected}}
(root / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
print('Pilot dataset ready', flush=True)