"""
http://www1.cs.columbia.edu/CAVE/software/softlib/coil-100.php

"Columbia Object Image Library (COIL-100),"
    S. A. Nene, S. K. Nayar and H. Murase, 
        Technical Report CUCS-006-96, February 1996. 

"""

import os, cPickle
import Image, numpy
from pylearn.datasets.config import data_root # config

from .memo import memo

def filenames():
    root = os.path.join(data_root(), 'COIL-100', 'coil-100', )
    for filename in os.listdir(root):
        yield filename, os.path.join(root,filename )

def filenameidx_imgidx(filename):
    if filename.startswith("obj"):
        obj_idx = int(filename[3:filename.index("_")])
        img_idx = int(filename[filename.index("_")+2:filename.index(".")])
        return obj_idx, img_idx
    else:
        raise ValueError(filename)

_32x32grey_path = os.path.join(data_root(), "COIL-100", "dct_32x32_grey.pkl")
_32x32grey_header = "Dictionary of COIL-100 dataset at 32x32 resolution, greyscale"
def build_32x32_grey():
    f = file(_32x32grey_path, "w")
    cPickle.dump(_32x32grey_header, f, protocol=cPickle.HIGHEST_PROTOCOL)

    dct = {}
    for filename, fullname in filenames():
        if filename.startswith('obj'):
            obj_idx, img_idx = filenameidx_imgidx(filename)
            img = numpy.asarray(Image.open(fullname))
            dct.setdefault(obj_idx, {})[img_idx] = img.mean(axis=2)[::4,::4]
    rval = numpy.empty((100, 72, 32, 32), dtype='float32')
    rval[...] = -1
    for obj_id, dd in dct.iteritems():
        for img_id, v in dd.iteritems():
            rval[obj_id, img_id, :, :] = v
    assert numpy.all(rval >= 0.0)

    cPickle.dump(rval, f, protocol=cPickle.HIGHEST_PROTOCOL)
    f.close()

@memo
def get_32x32_grey():
    f = file(_path_32x32_grey)
    if _32x32grey_header != cPickle.load(f):
        raise ValueError('wrong pickle file')
    rval = cPickle.load(f)
    f.close()
    return rval