import tensorflow as tf
import numpy as np
import skluc.mldatasets as dataset

IMAGE_SIZE = 24

val_size = 5000
cifar = dataset.Cifar10Dataset(validation_size=val_size)
cifar.load()
cifar.normalize()
cifar.to_one_hot()
cifar.data_astype(np.float32)
cifar.labels_astype(np.float32)
X_train, Y_train = cifar.train
X_test, Y_test = cifar.test
X_val, Y_val = cifar.validation


def distorded_inputs(image_tensor):
    height = IMAGE_SIZE

    width = IMAGE_SIZE

    distorted_image = tf.random_crop(image_tensor, [height, width, 3])
    distorted_image = tf.image.random_flip_left_right(distorted_image)
    distorted_image = tf.image.random_brightness(distorted_image,
                                                 max_delta=63)
    distorted_image = tf.image.random_contrast(distorted_image,
                                               lower=0.2, upper=1.8)
    float_image = tf.image.per_image_standardization(distorted_image)
    return float_image


# todo terminer ce programme
if __name__ == '__main__':
    SIGMA = 5.0
    print("Sigma = {}".format(SIGMA))

    with tf.Graph().as_default():
        output_dim = Y_train.shape[1]
        input_dim = X_train.shape[1:]

        x_image = tf.placeholder(tf.float32, shape=[None, *input_dim], name="x_image")
        y_ = tf.placeholder(tf.float32, shape=[None, output_dim], name="labels")
        tf.summary.image("cifarimage", x_image, max_outputs=10)

        dist_x_images = distorded_inputs(x_image)
        tf.summary.image("cifarimagedistorded", dist_x_images, max_outputs=10)

        # out = fast_food(x_image, SIGMA)

        merged_summary = tf.summary.merge_all()

        init = tf.global_variables_initializer()
        sess = tf.Session()

        summary_writer = tf.summary.FileWriter("cifar")
        feed_dict = {x_image: X_train[:10], y_: Y_train[:10]}
        summary = sess.run([merged_summary], feed_dict=feed_dict)
        summary_writer.add_summary(summary[0])
        summary_writer.close()