import numpy as np
import os
from os import listdir
from scipy import ndimage, misc
from keras.utils import to_categorical
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle

dataset_basepath = 'dataset/'

def process_data(path):
    try:
        X = np.load('data/X.npy')
        Y = np.load('data/Y.npy')
        print("Data already processed. Loading now.")
    except:
        print("Need to process data...")
        labels = sorted(listdir(path))
        X = []
        Y = []
        for i in labels:
            images_path = '{}{}'.format(path ,i)
            images = listdir(images_path)
            for img in images:
                image_path = '{}{}/{}'.format(path, i, img)
                image_data = ndimage.imread(image_path, mode="RGB")
                image_resized = misc.imresize(image_data, (64, 64))
                X.append(image_resized)
                Y.append(i)
        # Normalize the X data to zero mean
        X = np.array(X)/255
        # Convert labels to one-hot encoded vector
        Y = to_categorical(Y)

        if not os.path.exists('data/'):
            os.makedirs('data/')
        np.save('data/X.npy', X)
        np.save('data/Y.npy', Y)
    return X, Y

X, Y = process_data(dataset_basepath)
X, Y = shuffle(X, Y, random_state=0)

X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.25, random_state=42)