import cPickle
import itertools
import os
import numpy as np
from PIL import Image

# download the CIFAR-10 python dataset from http://www.cs.toronto.edu/~kriz/cifar.html
# it expands into a cifar-10-batches-py folder, which should be moved to this directory.

fo = open("./cifar-10-batches-py/data_batch_1", "rb")
d = cPickle.load(fo)
fo.close()

size = (32, 32)

cats = []
aps = []

for i, (data, label) in enumerate(itertools.izip(d["data"], d["labels"])):
    if i % 100 == 0:
        print i
    cat = label == 3
    airplane = label == 0
    if not cat and not airplane:
        continue
    if cat and len(cats) >= 100:
        continue
    if airplane and len(aps) >= 100:
        continue
    im = Image.new("RGB", size, "white")
    pixels = [(data[j], data[j+1024], data[j+2048]) for j in xrange(size[0]*size[1])]
    im.putdata(pixels)
    
    im = im.resize((8,8))
    im = im.convert("L")

    if cat:
        im.save("./cats/%d.png" % i)
    elif airplane:
        im.save("./airplanes/%d.png" % i)

    d = im.getdata()

    rkt = "'(" + " ".join("%d" % x for x in im.getdata()) + ")"
    if cat:
        cats.append(rkt)
    elif airplane:
        aps.append(rkt)


out = open("data.rkt", 'w')
out.write("#lang racket\n(provide cats airplanes)\n")
out.write("; this file is auto-generated by parse-data.py\n")

out.write("(define cats\n  (list\n")
for cat in cats:
    out.write("    %s\n" % cat)
out.write("  ))\n")

out.write("(define airplanes\n  (list\n")
for ap in aps:
    out.write("    %s\n" % ap)
out.write("  ))\n")

out.close()