"""Prepare reproducible CIFAR-10 cat/dog splits for the KNP example (numpy only)."""
import argparse
import hashlib
import json
from pathlib import Path
import tarfile
import urllib.request

import numpy as np

URL = 'https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz'
MD5 = 'c32a1d4ab5d03f1284b67883e8d87530'


def encode(rgb):
    # Fixed, untrained preprocessing. Separate brightness and darkness channels.
    gray = ((rgb[:, 0].astype(np.uint32) * 299 + rgb[:, 1].astype(np.uint32) * 587
             + rgb[:, 2].astype(np.uint32) * 114 + 500) // 1000).astype(np.uint8)
    return np.concatenate([gray, 255 - gray], axis=1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('output', type=Path)
    a = ap.parse_args()
    a.output.mkdir(parents=True, exist_ok=True)
    archive = a.output / 'cifar-10-binary.tar.gz'
    if not archive.exists():
        urllib.request.urlretrieve(URL, archive)
    assert hashlib.md5(archive.read_bytes()).hexdigest() == MD5, 'Archive checksum mismatch'
    with tarfile.open(archive) as tar:
        def read(names):
            data = np.concatenate([np.frombuffer(tar.extractfile('cifar-10-batches-bin/' + name).read(),
                                                  dtype=np.uint8).reshape(-1, 3073) for name in names])
            ids = np.flatnonzero(np.isin(data[:, 0], [3, 5]))
            return data[ids, 1:].reshape(-1, 3, 1024), (data[ids, 0] == 5).astype(np.uint8), ids
        rgb, labels, ids = read([f'data_batch_{i}.bin' for i in range(1, 6)])
        test_rgb, test_labels, test_ids = read(['test_batch.bin'])
    rng = np.random.default_rng(20261014)
    train, val = [], []
    for label in (0, 1):
        indices = rng.permutation(np.flatnonzero(labels == label))
        train.extend(indices[:4000])
        val.extend(indices[4000:])
    train = rng.permutation(train)
    val = rng.permutation(val)
    test = np.random.default_rng(20261015).permutation(len(test_labels))
    for name, x, y, source in [('train', rgb[train], labels[train], ids[train]),
                               ('val', rgb[val], labels[val], ids[val]),
                               ('test', test_rgb[test], test_labels[test], test_ids[test])]:
        encode(x).tofile(a.output / f'{name}.bin')
        np.savetxt(a.output / f'{name}.target', y, fmt='%d')
        np.savez_compressed(a.output / f'{name}-original.npz', rgb=x, labels=y, source_indices=source)
    for name, indices, inference in [('smoke', train[:100], val[:100]), ('validation', train, val)]:
        encode(np.concatenate([rgb[indices], rgb[inference]])).tofile(a.output / f'{name}.bin')
        np.savetxt(a.output / f'{name}.target', np.concatenate([labels[indices], labels[inference]]), fmt='%d')
    manifest = {'source': URL, 'md5': MD5, 'labels': {'0': 'cat', '1': 'dog'},
                'seed': 20261014, 'train': 8000, 'validation': 2000, 'test': 2000,
                'encoding': '32x32 grayscale integer BT.601; concatenate intensity and 255-intensity',
                'features': 2048, 'train_validation_disjoint': bool(not set(train) & set(val))}
    (a.output / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
    print(json.dumps(manifest), flush=True)


if __name__ == '__main__':
    main()