import pylab
from scipy import random, zeros, linspace, dot
import numpy, scipy
from pybrain.structure.modules import KohonenMap
from pybrain.datasets.unsupervised import UnsupervisedDataSet
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import itertools
import colorsys
import os

KOHONEN_SIZE = 3
KOHONEN_DIM = 8

som = KohonenMap(KOHONEN_DIM, KOHONEN_SIZE)

pylab.ion()

img = mpimg.imread('flowers.png')

IMAGE_WIDTH = len(img)
IMAGE_HEIGHT = len(img[0])

newimg = zeros((IMAGE_WIDTH, IMAGE_HEIGHT, KOHONEN_DIM))



image = mpimg.imread("flowers.png")
r, g, b = image[:,:,0], image[:,:,1], image[:,:,2]
h, s, v = r, g, b

for i in range(IMAGE_WIDTH):
  for j in range(IMAGE_HEIGHT):
    h[i,j] = colorsys.rgb_to_hsv(r[i,j], g[i,j], b[i,j])[0]
    s[i,j] = colorsys.rgb_to_hsv(r[i,j], g[i,j], b[i,j])[1]
    v[i,j] = colorsys.rgb_to_hsv(r[i,j], g[i,j], b[i,j])[2]
    if (colorsys.rgb_to_hsv(r[i,j], g[i,j], b[i,j])[2]*255.0 < 255): v[i,j] = 1 - v[i,j];

img[:,:,0] = v
img[:,:,1] = v
img[:,:,2] = v
'''
pylab.imshow(img)
pylab.draw()

'''
NEIGHBORING_RANGE = 1

for i in range(IMAGE_WIDTH):
  for j in range(IMAGE_HEIGHT):
    newimg[i,j,:3] = img[i,j,:3]
    add_feat = zeros((3))
    div = 0
    for k in range(i-NEIGHBORING_RANGE if i-NEIGHBORING_RANGE>0 else 0, i+NEIGHBORING_RANGE+1 if i+NEIGHBORING_RANGE+1<IMAGE_WIDTH else IMAGE_WIDTH-1):
      for m in range(j-NEIGHBORING_RANGE if j-NEIGHBORING_RANGE>0 else 0, j+NEIGHBORING_RANGE+1 if j+NEIGHBORING_RANGE+1<IMAGE_HEIGHT else IMAGE_HEIGHT-1):
          if not (k==i and m==j):
            div+=1
            for layer in range(3):
                add_feat[layer] += img[k,m,layer]
    add_feat = add_feat/[div]
    newimg[i,j,3:6] = add_feat
    newimg[i,j,6:] = [0.004*i,0.004*j]

dataset = newimg.reshape(IMAGE_WIDTH*IMAGE_HEIGHT, KOHONEN_DIM)
order = [i for i in range(len(dataset))]
#scipy.random.shuffle(order)
ds = UnsupervisedDataSet(KOHONEN_DIM)
for sample in dataset:
  ds.addSample(sample[:KOHONEN_DIM])

print("start learning")
for i in range(1):
  it = 0
  for n in order:
    point = dataset[n]
    # one forward and one backward (training) pass
    som.activate(point[:KOHONEN_DIM])
    som.backward()
    if it % 10000 == 0:
      print (it)
    it+=1

print("activating")
res = som.activateOnDataset(ds)
im = zeros((IMAGE_WIDTH, IMAGE_HEIGHT, 3))

colors = list(itertools.permutations([i*51 for i in range(6)], 3))

'''
#different colors
for i in range(IMAGE_WIDTH):
  for j in range(IMAGE_HEIGHT):
    n = res[i*IMAGE_WIDTH + j]
    im[i,j,:] = colors[int(n[0]*KOHONEN_SIZE + n[1])]

pylab.imshow(im)
pylab.draw()
'''

#avarage colors
im = zeros((IMAGE_WIDTH, IMAGE_HEIGHT, 3))
for i in range(IMAGE_WIDTH):
  for j in range(IMAGE_HEIGHT):
    n = res[i*IMAGE_WIDTH + j]
    im[i,j,:] = som.neurons[int(n[0]), int(n[1]), :3]

pylab.imshow(im)
pylab.draw()