#!/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

# Set GPU as only available physical device
print("Set GPU as only available device")
my_devices = tf.config.list_physical_devices(device_type='GPU')
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. 

# 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 = 1000
for x in range(repetitions):
    batch_images, batch_labels = next(it)