#!/bin/env python3
# https://www.tensorflow.org/api_docs/python/tf/keras/preprocessing/image/ImageDataGenerator
# https://stepup.ai/train_data_augmentation_keras/
import os
import tensorflow as tf
from tensorflow.keras.datasets import cifar10
from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.utils import to_categorical
import matplotlib.pyplot as plt

# Helper function to inspect the first images in a dataset
def visualize_data(images, categories, class_names, file_name):
    fig = plt.figure(figsize=(14, 6))
    fig.patch.set_facecolor('white')
    for i in range(3 * 7):
        plt.subplot(3, 7, i+1)
        plt.xticks([])
        plt.yticks([])
        plt.imshow(images[i])
        class_index = categories[i].argmax()
        plt.xlabel(class_names[class_index])
    fig.savefig(file_name)


# Set CPU as only available physical device (avoid GPU)
print("Set CPU as only available device")
my_devices = tf.config.list_physical_devices(device_type='CPU')
tf.config.set_visible_devices(devices= my_devices)

# CIFAR-10 Dataset
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
num_classes = len(class_names)

(x_train, y_train), (x_test, y_test) = cifar10.load_data()

x_train = x_train / 255.0
y_train = to_categorical(y_train, num_classes)

x_test = x_test / 255.0
y_test = to_categorical(y_test, num_classes)

# The above code first downloads the dataset. The included preprocessing rescales the images into the range
# between [0, 1] and converts the label from the class index (integers 0 to 10) to a one-hot encoded
# categorical vector. 

# Show the first images of the training set
visualize_data(x_train, y_train, class_names, 'train_samples.png')


# Specification of the augmentation parameters: 
#   width shift: Randomly shift the image left and right by 3 pixels
#   height shift: Randomly shift the image up and down by 3 pixels
#   horizontal flip: Randomly flip the image horizontally.

width_shift = 3/32
height_shift = 3/32
flip = True

datagen = ImageDataGenerator(
    horizontal_flip=flip,
    width_shift_range=width_shift,
    height_shift_range=height_shift,
    )
datagen.fit(x_train)

# output directory
path = "./augmented_data"
try:
    os.mkdir(path)
except OSError:
    print ("Creation of the directory %s failed" % path)
else:
    print ("Successfully created the directory %s " % path)

my_batch_size=32
it = datagen.flow(x_train, y_train, shuffle=False, batch_size=my_batch_size,
        save_to_dir=path, save_prefix='artemisa_ex2')

# If we want to iterate and create more batches of images

# num_images_generated = batch_size * repetitions
# repetitions = 100
# for x in range(repetitions):
batch_images, batch_labels = next(it)

# Show samples augmented data
visualize_data(batch_images, batch_labels, class_names, 'augmented_samples.png')