import os, numpy
import theano

from pylearn.datasets.image_patches import  (
        olshausen_field_1996_whitened_images,
        extract_random_patches)

from .protocol import TensorFnDataset # protocol.py __init__.py
from .memo import memo

@memo
def get_dataset(N,R,C,dtype):
    seed=98234
    rng = numpy.random.RandomState(seed)
    img_stack = olshausen_field_1996_whitened_images()
    patch_stack = extract_random_patches(img_stack, N,R,C,rng)
    return patch_stack.astype(dtype).reshape((N,(R*C)))

def image_patches(s_idx, dims,
        split='train', dtype=theano.config.floatX, rasterized=False):
    N,R,C=dims

    if split != 'train':
        raise NotImplementedError('train/test/valid splits for randomly sampled image patches?')

    if not rasterized:
        raise NotImplementedError()

    op = TensorFnDataset(dtype, bcast=(False,), fn=(get_dataset, (N,R,C,dtype,)), single_shape=(R*C,))
    x = op(s_idx%N)
    if x.ndim == 1:
        if not rasterized:
            x = x.reshape((20,20))
    elif x.ndim == 2:
        if not rasterized:
            x = x.reshape((x.shape[0], 20,20))
    else:
        assert False, 'what happened?'

    return x