import PIL.Image as Image
import PIL.ImageDraw as ImageDraw
import colorsys
import random

from pybrain.tools.shortcuts import buildNetwork
from pybrain.datasets import SupervisedDataSet
from pybrain.supervised.trainers import BackpropTrainer
import pybrain.structure.modules

import scipy.misc
import scipy.fftpack
import scipy


wood = Image.open('wood.png')
notwood = Image.open('notwood.png')

def get_sample(image, x, y):
	'''Preprocess pixel values to neural net inputs.
	Converts middle pixel to HSL.
	Takes 32x32 tile and average of its FFT.
	'''
	r,g,b = image.getpixel((x,y))[:3]
	hls = colorsys.rgb_to_hls(r/255., g/255., b/255.)
	
	x1 = max(0, x - 8)
	y1 = max(0, y - 8)
	x2 = min(image.size[0], x + 8) # PIL skips the end coordinate
	y2 = min(image.size[1], y + 8)
	tile = image.crop((x1,y1,x2,y2)).convert('L')
	
#	arr = scipy.misc.fromimage(tile) / 255.
#	fft = abs(scipy.fftpack.fft2(arr))
#	fft = fft[1:-1, 1:-1]
#	avg_fft = fft.sum() / (fft.shape[0] * fft.shape[1])
	
	l_min, l_max = tile.getextrema()
	
	return hls + (l_min / 255., l_max / 255.)

def addimage(image, output):
	data = image.getdata()
	results = []
	for y in range(0, image.size[1], 1):
		for x in range(0, image.size[0], 1):
			sample = get_sample(image, x, y)
			results.append((sample, output))
	return results

samples = addimage(wood, (1,))
samples += addimage(notwood, (-1,))

input_count = len(samples[0][0])

random.shuffle(samples)
dataset = SupervisedDataSet(input_count,1)
file1 = open('data1.txt', 'w')
file2 = open('data2.txt', 'w')
for sample, output in samples:
	if output[0] > 0:
		file1.write(' '.join(map(str,sample)) + '\n')
	else:
		file2.write(' '.join(map(str,sample)) + '\n')
	dataset.addSample(sample, output)

print "Got", len(dataset), "samples"

net = buildNetwork(input_count, input_count, 1) #, outclass = pybrain.structure.modules.TanhLayer)
trainer = BackpropTrainer(net, dataset)
print "Training 1", trainer.train()
print "Training 2", trainer.train()
print "Training 3", trainer.train()
#print "Training 4", trainer.train()
#print "Training 5", trainer.train()

dataset = SupervisedDataSet(input_count,1)
test = Image.open('wms_nuuksio2.png')
data = test.getdata()
result = []
step = 4
for y in range(0, test.size[1], step):
	for x in range(0, test.size[0], step):
		sample = get_sample(test, x, y)
		dataset.addSample(sample, 0)

output = net.activateOnDataset(dataset)

draw = ImageDraw.Draw(test)
i = 0
for y in range(0, test.size[1], step):
	for x in range(0, test.size[0], step):
		value = output[i][0]
		draw.rectangle((x,y,x + step, y + step), fill = (int(value * 255), 0, 0))
		i += 1

test.save('output.png')