"""Small pilot using JPEG previews of uoft-cs/cifar10, not canonical CIFAR bytes."""
from concurrent.futures import ThreadPoolExecutor
import hashlib
from io import BytesIO
import json
from pathlib import Path
import time
import urllib.parse
import urllib.request
import urllib.error

import numpy as np
from PIL import Image

from prepare_data import encode

root = Path('/home/int/run/cats-dogs/data')
cache = root / 'preview-cache'
cache.mkdir(exist_ok=True)


def request(url):
    for attempt in range(8):
        try:
            with urllib.request.urlopen(url, timeout=45) as response:
                return response.read()
        except Exception as error:
            if attempt == 7:
                raise
            delay = min(45, 3 * (attempt + 1))
            if isinstance(error, urllib.error.HTTPError) and error.code == 429:
                delay = max(delay, int(error.headers.get('Retry-After', '15')))
            time.sleep(delay)


def rows(args):
    split, offset = args
    cached = cache / f'rows-{split}-{offset}.json'
    if cached.exists():
        return split, json.loads(cached.read_text())['rows']
    params = urllib.parse.urlencode(dict(dataset='uoft-cs/cifar10', config='plain_text',
                                        split=split, offset=offset, length=100))
    response = request('https://datasets-server.huggingface.co/rows?' + params)
    cached.write_bytes(response)
    time.sleep(0.5)
    return split, json.loads(response)['rows']


catalog = {'train': [], 'test': []}
tasks = [('train', i) for i in range(0, 6000, 100)] + [('test', i) for i in range(0, 3000, 100)]
with ThreadPoolExecutor(max_workers=2) as pool:
    for split, batch in pool.map(rows, tasks):
        catalog[split].extend(r for r in batch if r['row']['label'] in (3, 5))
selected = {'train': [], 'val': [], 'test': []}
rng = np.random.default_rng(20261014)
for label in (3, 5):
    candidates = sorted([r for r in catalog['train'] if r['row']['label'] == label], key=lambda r: r['row_idx'])
    candidates = [candidates[i] for i in rng.permutation(len(candidates))]
    assert len(candidates) >= 500
    selected['train'].extend(candidates[:400])
    selected['val'].extend(candidates[400:500])
    candidates = sorted([r for r in catalog['test'] if r['row']['label'] == label], key=lambda r: r['row_idx'])
    assert len(candidates) >= 200
    selected['test'].extend(candidates[:200])
for split in selected:
    selected[split] = [selected[split][i] for i in rng.permutation(len(selected[split]))]


def image(args):
    split, row = args
    path = cache / f'{split}-{row["row_idx"]}.jpg'
    if not path.exists():
        path.write_bytes(request(row['row']['img']['src']))
    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)


arrays, labels = {}, {}
provenance = {}
with ThreadPoolExecutor(max_workers=4) as pool:
    for split in selected:
        original_split = 'test' if split == 'test' else 'train'
        arrays[split] = np.stack(list(pool.map(image, [(original_split, row) for row in selected[split]])))
        labels[split] = np.array([int(row['row']['label'] == 5) for row in selected[split]], dtype=np.uint8)
        provenance[split] = [row['row_idx'] for row in selected[split]]
        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],
                            source_indices=provenance[split])
        print(split, len(labels[split]), '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(provenance['train']) & set(provenance['val'])
manifest = {'source': 'https://huggingface.co/datasets/uoft-cs/cifar10',
            'variant': 'JPEG previews from datasets-server API; 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_indices': provenance, '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)