"""
Routines to load/access Olshausen's image_patches
"""

import os
import numpy

import scipy.io  #for loadmat

from .config import data_root
from .dataset import Dataset

dirpath = os.path.join(data_root(), 'image_patches','olshausen','smaller_patches')

paths = {'20by20_whiten_01': ('natural_images_patches_whiten.amat',(20,20)),
         '12by12_whiten_01': ('natural_images_patches_whiten_12_by_12_0_1.amat',(12,12))}

def load_dataset(ntrain=70000, nvalid=15000, ntest=15000, variant='20by20_whiten_01'):
    # This is implementation loads files which are way too big (storing bytes in double), and
    # stored in ascii!??, and which appear to be un-documented variants on the original
    # raw/whitened
    # data.
    # I'd like to deprecate this function.
    # -JB Aug 2010
    print >> sys.stderr, "WARNING: pylearn.datasets.image_patches.load_dataset is badly documented and does some weird stuff... could someone who uses this function do something about it?"
    
    ndata = 100000

    if not paths.get(variant, None):
        raise ValueError('Unknown image_patches variant: %s' % variant)
    if ntrain + nvalid + ntest < ndata:
        raise ValueError('ntrain + nvalid + ntest must be smaller than %i' %ndata)

    fname = os.path.join(dirpath, paths[variant][0])
    data = numpy.loadtxt(fname)
    x = data[:,:-1]
    y = data[:,-1]

    perm = numpy.random.permutation(ndata)
 
    rval = Dataset()
    rval.train = Dataset.Obj(x = x[perm[:ntrain],:], y = y[perm[:ntrain]])
    rval.valid = Dataset.Obj(x = x[perm[ntrain:ntrain+nvalid],:],
                             y = y[perm[ntrain:ntrain+nvalid]])
    rval.test  = Dataset.Obj(x = x[perm[:-ntest],:],
                             y = y[perm[:-ntest]])

    rval.n_classes = 10
    rval.img_shape = paths[variant][1]

    return rval


#TODO: a little function to render a tiling of example images to an image using PIL

#TODO: a pca_load_dataset function that loads the data, as projected onto principle components

def olshausen_field_1996_whitened_images(path=None):
    """Returns a (10,512,512) ndarray containing 10 whitened images.
    
    The images are in floating point.
    Whitening was done by the paper authors, with band-pass whitening I think.
    """
    if path is None:
        path=os.path.join(data_root(), 'image_patches', 'olshausen',
                'original', 'IMAGES.mat')
    images = scipy.io.loadmat(path)['IMAGES']
    assert images.shape == (512,512,10)
    return images.astype('float32').transpose([2,0,1])

def olshausen_field_1996_raw_images(path=None):
    """Returns a (10,512,512) ndarray containing 10 images.
    
    The images are in floating point.
    """
    if path is None:
        path=os.path.join(data_root(), 'image_patches', 'olshausen',
                'original', 'IMAGES_RAW.mat')
    images = scipy.io.loadmat(path)['IMAGES_RAW']
    assert images.shape == (512,512,10)
    return images.astype('float32').transpose([2,0,1])

def extract_random_patches(img_stack, N, R,C, rng):
    """Return subimages from the img_stack

    :param img_stack: a 3D[4D] ndarray (n_images, rows, cols,[channels]) or a list of images.
    :param N: number of patches to extract
    :param R: number of rows in patch
    :param C: number of cols in patch
    :param rng: numpy RandomState
    
    Sub-image regions are chosen at random from the img_stack with uniform probability, and
    then within each image with uniform probability across the image.  Patches from a larger
    image in the stack therefore would be sampled less frequently than patches from a smaller
    image in the stack.

    :hint:
        To use ZCA whitening, extract patches from the raw data, and pass it to
        preprocessing.pca.zca_whitening.
    """

    rval = numpy.empty((N,R,C)+img_stack.shape[3:], dtype=img_stack[0].dtype)

    L = len(img_stack)
    img_idxlist = rng.randint(L,size=N)
    offsets_R = rng.randint(img_stack.shape[1]-R+1, size=N)
    offsets_C = rng.randint(img_stack.shape[2]-C+1, size=N)

    for n, (l,r,c) in enumerate(zip(img_idxlist, offsets_R, offsets_C)):
        #img_idx = rng.randint(L)
        #offset_R = rng.randint(img_n.shape[0]-R+1)
        #offset_C = rng.randint(img_n.shape[1]-C+1)
        #img_n = img_stack[l]
        #rval[n] = img_n[offset_R:offset_R+R,offset_C:offset_C+C]
        rval[n] = img_stack[l,r:r+R,c:c+C]
    return rval