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_landsat.png')

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)
		samples.append(sample)
print "done"
		
input_count = len(samples[0])

samples.append([0.50] * input_count)

gmm = GMM(5, 'diag')

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

print "Collecting testing samples...",
test = Image.open('wms_landsat.png')
output = Image.new('RGB', test.size)
draw = ImageDraw.Draw(output)
step = 8
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 = gmm.predict(numpy.array(row))
	
	print results
	
	i = 0
	for x in range(0, test.size[0], step):
		value = results[i]
		
		draw.rectangle((x,y,x + step, y + step), fill = (0, ((value * 9999) // 256) % 256, (value * 9999) % 256))
		i += 1
print "done"
		
output.save('output.png')