'''
Created by Tony Silvestre to prepare images for use from a Kaggle Where's Waldo dataset
'''
import os
import numpy as np
from matplotlib import pyplot as plt
import math
import cv2
from skimage import color, exposure, transform

'''
Defines some basic preprocessing for the images
'''
def preprocess(img):
    # Histogram normalization in v channel
    hsv = color.rgb2hsv(img)
    hsv[:, :, 2] = exposure.equalize_hist(hsv[:, :, 2])
    img = color.hsv2rgb(hsv)

    img = img/255                           # Scaling images down to values of 0-255
    img = np.rollaxis(img, -1)              # rolls colour axis to 0 

    return img

def gen_data(w_path, n_w_path):
    waldo_file_list = os.listdir(os.path.join(w_path))
    total_w = len(waldo_file_list)
    not_waldo_file_list = os.listdir(os.path.join(n_w_path))
    total_nw = len(not_waldo_file_list)
    imgs_raw = []           # Images
    imgs_lbl = []           # Image labels

    i = 0
    for image_name in waldo_file_list:
        pic = cv2.imread(os.path.join(w_path, image_name))              # NOTE: cv2.imread() returns a numpy array in BGR not RGB
        pic = preprocess(pic)
        
        imgs_raw.append(pic)
        imgs_lbl.append(1)                                              # Value of 1 as Waldo is present in the image

        print('Completed: {0}/{1} Waldo images'.format(i+1, total_w))
        i += 1

    i = 0
    for image_name in not_waldo_file_list:
        pic = cv2.imread(os.path.join(n_w_path, image_name))
        pic = preprocess(pic)
        
        imgs_raw.append(pic)
        imgs_lbl.append(0)
        
        print('Completed: {0}/{1} non-Waldo images'.format(i+1, total_nw))
        i += 1

    ## Randomise and split data into training and test sets
    # Code was modified from code written by: Kyle O'Brien (medium.com/@kylepob61392)
    n_images = len(imgs_raw)
    TRAIN_TEST_SPLIT = 0.75

    # Split at the given index
    split_index = int(TRAIN_TEST_SPLIT * n_images)
    shuffled_indices = np.random.permutation(n_images)
    train_indices = shuffled_indices[0:split_index]
    test_indices = shuffled_indices[split_index:]

    train_data = []
    train_lbl = []
    test_data = []
    test_lbl = []

    # Split the images and the labels
    for index in train_indices:
        train_data.append(imgs_raw[index])
        train_lbl.append(imgs_lbl[index])

    for index in test_indices:
        test_data.append(imgs_raw[index])
        test_lbl.append(imgs_lbl[index])

    # # Calculate what 30% of each set is
    # third_of_w = math.floor(0.3*total_w)
    # third_of_nw = math.floor(0.3*total_nw)
    
    # # Split data into training and test data (60%/30%)    
    # train_data = np.append(imgs_raw[(third_of_w+1):total_w], imgs_raw[(total_w + third_of_nw + 1):len(imgs_raw)-1], axis=0)
    # train_lbl = np.append(imgs_lbl[(third_of_w+1):total_w], imgs_lbl[(total_w + third_of_nw + 1):len(imgs_lbl)-1], axis=0)
    # # If axis not given, both arrays are flattened before being appended
    # test_data = np.append(imgs_raw[0:third_of_w], imgs_raw[total_w:(total_w + third_of_nw)], axis=0)
    # test_lbl = np.append(imgs_lbl[0:third_of_w], imgs_lbl[total_w:(total_w + third_of_nw)], axis=0)

    try:
        # Save the data as numpy files
        np.save('Waldo_train_data.npy', train_data)
        np.save('Waldo_train_lbl.npy', train_lbl)
        np.save('Waldo_test_data.npy', test_data)
        np.save('Waldo_test_lbl.npy', test_lbl)
        print("All data saved")
    except:
        print("ERROR: Data may not be completely saved")


if __name__ == "__main__":
    # Paths to the Waldo images
    waldo_path = 'waldo_data/64/waldo'
    n_waldo_path = 'waldo_data/64/notwaldo'

    gen_data(waldo_path, n_waldo_path)