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

import texture

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))
	
#	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])
	
	extremas = tile.getextrema()
	
	return hls + tuple([e[0]/255. for e in extremas] + [e[1]/255. for e in extremas] + [(e[1] - e[0])/255. for e in extremas]) + tuple(texture.get_sample(image, x, y))


wood = Image.open('wms_2465_2467_6027_6028.png')
mask = Image.open('mask.png').convert('RGB')

class Status:
	def __init__(self):
		self.status = 0
	
	def __call__(self, now, total):
		newstatus = int(10 * now / total) * 10
		if newstatus != self.status:
			self.status = newstatus
			print str(newstatus) + "%",

status = Status()
			
print "Collecting training samples..",
samples = []
for y in range(0, wood.size[1], 8):
	status(y, wood.size[1])
	for x in range(0, wood.size[0], 8):
		sample = get_sample(wood, x, y)
		maskvalue = mask.getpixel((x,y))
		if maskvalue == (0,255,0):
			output = 1
		elif maskvalue == (255,0,0):
			output = -1
		else:
			continue
		
		samples.append((sample, output))
print "done"
		
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()

print "Collecting testing samples...",
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):
	status(y, test.size[1])
	for x in range(0, test.size[0], step):
		sample = get_sample(test, x, y)
		dataset.addSample(sample, 0)
print "done"

print "Classifying test data...",
output = net.activateOnDataset(dataset)
print "done"

print "Drawing result image...",
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 * 128 + 128), 0, 0))
		i += 1
print "done"
		
test.save('output.png')