'''
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
import skimage as sk
import random

'''
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)

    return img

'''
Augments an image (horizontal reflection, hue, contrast, saturation adjustment) to expand on the dataset
Code was adapted from Medium.com/@thimblot (Thomas Himblot)
'''
def augment(img):
    # Randomly rotate the image
    random_degree = random.uniform(-25, 25)         # Random degree of rotation between 25% on the left and 25% on the right
    img = sk.transform.rotate(img, random_degree)

    # Add random noise to the image
    img = sk.util.random_noise(img)

    # Randomly (25% chance) flips the image horizontally (by inverting the image dimensions)
    if random.randint(0, 4) == 1:
        img = img[:, ::-1]

    # Randomly (25% chance) inverts the intensity range of the image
    if random.randint(0, 4) == 1:
        img = sk.util.invert(img)

    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

    w = 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)

        pic_roll = np.rollaxis(pic, -1)                                 # rolls colour axis to 0 
        imgs_raw.append(pic_roll)
        imgs_lbl.append(1)                                              # Value of 1 as Waldo is present in the image

        for j in range(0, 10):
            pic = augment(pic)
            pic_roll = np.rollaxis(pic, -1)                                 # rolls colour axis to 0 
            imgs_raw.append(pic_roll)
            imgs_lbl.append(1)                          # Value of 1 as Waldo is still present in the transformed image

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

    nw = 0
    for image_name in not_waldo_file_list:
        pic = cv2.imread(os.path.join(n_w_path, image_name))
        pic = preprocess(pic)
        
        pic = np.rollaxis(pic, -1)                                      # rolls colour axis to 0 
        imgs_raw.append(pic)
        imgs_lbl.append(0)
        
        print('Completed: {0}/{1} non-Waldo images'.format(nw+1, total_nw))
        if nw > 10*w:
            print("Non_Waldo files restricted")
            break
        else:
            nw += 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             # Amount of trainingdata as a percentage of the total

    # 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])

    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)