#!/usr/bin/env python

import matplotlib
#matplotlib.use('Agg')
import matplotlib.pyplot as P
import matplotlib.image as mpimg
import os, sys, glob
import numpy as N

class data(object):
    def __init__(self, path, width, height, sizeRetina):
        self.__width = width
        self.__height = height
        self.__size = sizeRetina
        tmp = str(path) + '/*.png'
        counter = 0
        for infile in glob.glob(tmp):
            img=mpimg.imread(str(infile))
            if  img.shape[0] == self.__height:
                counter += 1
        self.__data = N.zeros((counter, self.__height, self.__width))
        print self.__data.shape
        print "Number of images: %d" %(counter)
        self.__numOfData = counter
        counter = 0
        for infile in glob.glob(tmp):           
            print "processing: " + infile
            img=mpimg.imread(str(infile))
            if img.shape[0] == self.__height:
                self.__data[counter] = N.copy(img)#-N.mean(img))
                counter += 1
#        P.imshow(self.__data[0], cmap=P.cm.gray)
#        P.show()
        # compute safety margin
        margin = self.__size # N.ceil(self.__size/2.)
        self.__range_width = N.array([int(margin), int(self.__width-margin)])
        self.__range_height = N.array([int(margin), int(self.__height-margin)])

    def retina(self, n, w , h):
        offset = self.__size/2
        # subtract mean
        lo = w - offset
        hi = w + offset
        bo = h - offset
        top = h + offset
        # subtract mean
        tmp = self.__data[n,bo:top,lo:hi]
        return N.copy((tmp-N.mean(tmp)))
#        return N.copy(tmp)

    def getPatch(self):
        # choose random picture
        rnd = N.random.randint(0,self.__numOfData)
        # pic random retina location within picture
        w = N.random.randint(self.__range_width[0],self.__range_width[1] )
        h = N.random.randint(self.__range_height[0], self.__range_height[1])
#        print w , ' -- ' , h       
#        print rnd
        return N.copy(self.retina(rnd,w,h))

if __name__ == "__main__":
    print "Retina"
    myData = data('data', 481, 321, 16)
    tmp = myData.getPatch()
    P.matshow(tmp, cmap=P.cm.gray)
    P.colorbar()
    P.show()