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

from scikits.learn.gmm import GMM

import scipy.misc
import scipy.fftpack
import scipy

import numpy

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))
	
	extremas = tile.getextrema()
	
	x1 = max(0, x - 1)
	y1 = max(0, y - 1)
	x2 = min(image.size[0], x + 2) # PIL skips the end coordinate
	y2 = min(image.size[1], y + 2)
	tile = image.crop((x1,y1,x2,y2))
	extremas2 = 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] + [(e[1] - e[0])/255. for e in extremas2]) + 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_wood = []
samples_notwood = []
for y in range(0, wood.size[1], 4):
	status(y, wood.size[1])
	for x in range(0, wood.size[0], 4):
		maskvalue = mask.getpixel((x,y))
		if maskvalue == (0,255,0):
			sample = get_sample(wood, x, y)
			samples_wood.append(sample)
#		elif maskvalue == (255,0,0):
#			samples_notwood.append(sample)
print "done"
		
input_count = len(samples_wood[0])

samples_wood.append([0.50] * input_count)
#samples_notwood.append([0.50] * input_count)

gmm_wood = GMM(n_states = 5, n_dim = input_count)
#gmm_notwood = GMM(n_states = 10, n_dim = input_count)

print "Fitting GMM's:",
print "gmm_wood",
gmm_wood.fit(numpy.array(samples_wood))
#print "gmm_notwood",
#gmm_notwood.fit(numpy.array(samples_notwood))
print "done"

print "Collecting testing samples...",
test = Image.open('wms_nuuksio2.png')
output = Image.new('RGB', test.size)
draw = ImageDraw.Draw(output)
step = 4
for y in range(0, test.size[1], step):
	status(y, test.size[1])
	
	row = []
	for x in range(0, test.size[0], step):
		sample = get_sample(test, x, y)
		row.append(sample)
	
	results_wood = gmm_wood.score(numpy.array(row))
	
	i = 0
	for x in range(0, test.size[0], step):
		value_wood = results_wood[i]
		
		draw.rectangle((x,y,x + step, y + step), fill = (0, int(value_wood) + 128, 0))
		i += 1
print "done"
		
output.save('output.png')