#!/usr/bin/env python


from pylearn.io import filetensor as ft
import copy
import pygame
import time
import numpy as N

#from gimpfu import *


from PoivreSel import PoivreSel
from thick import Thick
from BruitGauss import BruitGauss
from DistorsionGauss import DistorsionGauss
from PermutPixel import PermutPixel
from gimp_script import GIMP1
from Rature import Rature
from contrast import Contrast
from local_elastic_distortions import LocalElasticDistorter
from slant import Slant
from Occlusion import Occlusion
from add_background_image import AddBackground
from affine_transform import AffineTransformation

###---------------------order of transformation module
MODULE_INSTANCES = [Slant(),Thick(),AffineTransformation(),LocalElasticDistorter(),GIMP1(), PermutPixel(),DistorsionGauss(), Rature(), Occlusion(),AddBackground(), BruitGauss(),PoivreSel(), Contrast()]

###---------------------complexity associated to each of them
complexity = [0.7,0.7,0.7,0.7,0.7,0.3,0.3,0.5,0.5,0.1,0.3,0.3,0.5]
#complexity = [0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.]



nbmodule = len(MODULE_INSTANCES)

datapath = '/data/lisa/data/nist/by_class/'
f = open(datapath+'lower/lower_train_data.ft')
d = ft.read(f)

d = d[0:1000,:]/255.0

pygame.surfarray.use_arraytype('numpy')

#pygame.display.init()
screen = pygame.Surface((4*(nbmodule+1)/2*32,2*(4*32+20)),depth=32)

anglcolorpalette=[(x,x,x) for x in xrange(0,256)]
#pygame.Surface.set_palette(anglcolorpalette)
#screen.set_palette(anglcolorpalette)

pygame.font.init()

for i in range(10000):
    a=d[i,:]
    b=N.asarray(N.reshape(a,(32,32)))
    c=N.asarray([N.reshape(a*255.0,(32,32))]*3).T
    new=pygame.surfarray.make_surface(c)
    new=pygame.transform.scale2x(new)
    new=pygame.transform.scale2x(new)
    #new.set_palette(anglcolorpalette)
    screen.blit(new,(0,0))
    
    offset = 4*32
    offset2 = 0
    ct = 0
    for j in MODULE_INSTANCES:
        #max dilation
             
        #random
        print j.get_settings_names(), j.regenerate_parameters(N.random.rand()*complexity[ct])

        b=j.transform_image(b)
        c=N.asarray([b*255]*3).T
        
        new=pygame.surfarray.make_surface(c)
        new=pygame.transform.scale2x(new)
        new=pygame.transform.scale2x(new)
        #new.set_palette(anglcolorpalette)
        screen.blit(new,(offset,offset2))
        font = pygame.font.SysFont('liberationserif',18)
        text = font.render(j.__module__,0,(255,255,255),(0,0,0))
        screen.blit(text,(offset,offset2+4*32))
        if ct == len(MODULE_INSTANCES)/2-1:
            offset = 0
            offset2 = 4*32+20
        else:
            offset += 4*32
        ct+=1
    pygame.image.save(screen,'/u/glorotxa/exemples/%s.BMP'%i)
    #raw_input('Press Enter')

#pygame.display.quit()