#!/usr/bin/python
# coding: utf-8

'''
Simple implementation of random affine transformations based on the Python 
Imaging Module affine transformations.


Author: Razvan Pascanu
'''

import numpy, Image



class AffineTransformation():
    def __init__( self, shape = (32,32), seed = None):
        self.shape = shape
        self.rng = numpy.random.RandomState(seed)

    def transform(self,NIST_image):
    
        im = Image.fromarray( \
                numpy.asarray(\
                       NIST_image.reshape(self.shape), dtype='uint8'))
        # generate random affine transformation
        # a point (x',y') of the new image corresponds to (x,y) of the old
        # image where : 
        #   x' = params[0]*x + params[1]*y + params[2]
        #   y' = params[3]*x + params[4]*y _ params[5]

        # the ranges are set manually as to look acceptable
        params = self.rng.uniform(size = 6) -.5
        params[2] *= 8.
        params[5] *= 8.
        params[0] = 1. + params[0]*0.4
        params[3] = 0. + params[3]*0.4
        params[1] = 0  + params[1]*0.4
        params[4] = 1  + params[4]*0.4

        print params
        nwim = im.transform( (32,32), Image.AFFINE, params)
        return numpy.asarray(nwim)



if __name__ =='__main__':
    print 'random test'
    
    from pylearn.io import filetensor as ft
    import pylab

    datapath = '/data/lisa/data/nist/by_class/'

    f = open(datapath+'digits/digits_train_data.ft')
    d = ft.read(f)
    f.close()


    transformer = AffineTransformation()
    id = numpy.random.randint(30)
    
    pylab.figure()
    pylab.imshow(d[id].reshape((32,32)))
    pylab.figure()
    pylab.imshow(transformer.transform(d[id]).reshape((32,32)))

    pylab.show()