"""
Created on Thu Oct 26 20:50:49 2017
Multiprocessing program of independent pool workers uses the Pillow 3
module to open PNG image files. The pool workers each individually
convert the PNG image file to a numpy array of shape (300, 300, 3) where
300 represents the width and height of the image and 3 represents color
format. The image's corresponding y_label is generated by doing keras'
to_categorical encoding that transforms integers into different binary
1,0 encoding of classes. These classes are circle, rectangle, triangle,
and square. The numpy array's total size of 2000 images per shape is
reduced by a half if it's numpy float 16 instead of numpy float 32.
@author: maggie
"""
from __future__ import print_function
from PIL import Image
import glob
import pickle
import numpy as np
from multiprocessing import Lock
from multiprocessing import Pool
from keras.utils import to_categorical

"""to make lock and queue storage global to all child workers in the Pool"""
def init(lock):
    global childs_lock
    childs_lock = lock

# this section of the code is for saving background colors
# opens the text file with the name of image and the rgb colors
def save_background_colors(image_path, file_txt):
    childs_lock.acquire()
    f = open(file_txt, 'r')
    img = Image.open(image_path)
    childs_lock.release()

    previous_line = " "
    float_numbers = []
    for current_line in f:
		if image_path in previous_line:
			y = [current_line.strip() for current_line in current_line.split(' ')]
			#remove string literals by casting float to string
			for numbers in y:
				c = float(numbers)
				float_numbers.append(c)
			print (image_path, y)
			background_tuple = tuple(float_numbers)
		previous_line = current_line
    f.close()
    np_img = np.array(img, dtype = [('img_info', np.float16)])

    return np_img['img_info'], background_tuple

def process_images(image_path, shape_path):
    shape_y = None
    if shape_path == "crop_images/valid_circle/":
        shape_y = 0
    if shape_path == "crop_images/valid_rectangle/":
        shape_y = 1
    if shape_path == "crop_images/valid_triangle/":
        shape_y = 2
    if shape_path == "crop_images/valid_square/":
        shape_y = 3

    ylabel = to_categorical(shape_y, num_classes = 4)
    ylabel = np.reshape(ylabel, (4))

    childs_lock.acquire()
    img = Image.open(image_path)
    childs_lock.release()
    np_img = np.array(img, dtype = [('img_info', np.float16)])
    #img = img.resize((200, 200), Image.ANTIALIAS) #ANTIALIAS reserves quality
    # to check that all image input are the same shape
    '''width, height = np_img.shape[0], np_img.shape[1]
    if width == 300 or height == 300:
		print ("rm ", image_path)'''
    img.close()
    return np_img['img_info'], ylabel

# global storage variable for both main and pool of workers
pickle_file = 'crop_images/valid_circle.pkl'
# create empty pickle_file first then append to file
output = open (pickle_file, 'wb')
output.close()

def result(data):
    output = open (pickle_file, 'ab')
    #print ("in pickle file: " , pickle_file)
    pickle.dump(data, output, pickle.HIGHEST_PROTOCOL)
    output.close()

if __name__ == '__main__':

    shape_path = "crop_images/valid_circle/"
    lock = Lock()
    p = Pool(processes = 4, initargs = (lock, ), initializer = init)

    for image_path in glob.glob(shape_path + "*png"):
        p.apply_async(process_images, (image_path, shape_path), callback = result)
    p.close() # no more tasks
    p.join() # wrap up current tasks

    file_txt = "background_circle.txt"
    for image_path in glob.glob(shape_path + "*png"):
        p.apply_async(save_background_colors, (image_path, file_txt))
    p.close() # no more tasks
    p.join() # wrap up current tasks