import tensorflow as tf
import numpy as np
from tensorflow.keras.callbacks import CSVLogger
from tensorflow.keras.preprocessing.image import ImageDataGenerator
mnist = tf.keras.datasets.fashion_mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
x_test = x_test.reshape(x_test.shape[0], 28, 28, 1)
x_train, x_test = x_train / 255.0, x_test / 255.0

#y_train = tf.keras.utils.to_categorical(y_train)
y_test = tf.keras.utils.to_categorical(y_test)

def get_random_sample(a, b, number_of_samples=10):
    x = []
    y = []
    for category_number in range(0,10):
        # get all samples of a category
        train_data_category = a[b==category_number]
        # pick a number of random samples from the category
        train_data_category = train_data_category[np.random.randint(train_data_category.shape[0], 
                                                                    size=number_of_samples), :]
        x.extend(train_data_category)
        y.append([category_number]*number_of_samples)

    return np.asarray(x).reshape(-1, 28, 28, 1), np.asarray(y).reshape(10*number_of_samples,1)

for i in ['1']:

    model = tf.keras.Sequential()

    model.add(tf.keras.layers.Conv2D(filters = 32, kernel_size = (3, 3), activation='relu',
                 input_shape = (28, 28, 1), padding='same'))
    model.add(tf.keras.layers.Conv2D(filters = 32, kernel_size = (2, 2), activation='relu', padding = 'same'))
    model.add(tf.keras.layers.MaxPool2D(strides=(2,2)))

    model.add(tf.keras.layers.Conv2D(filters = 64, kernel_size = (3, 3), activation='relu', padding='same'))
    model.add(tf.keras.layers.Conv2D(filters = 64, kernel_size = (3, 3), activation='relu', padding='same'))
    model.add(tf.keras.layers.MaxPool2D(strides=(2,2)))

    model.add(tf.keras.layers.Flatten())

    model.add(tf.keras.layers.Dense(256, activation='relu'))
    model.add(tf.keras.layers.Dropout(0.2))
    model.add(tf.keras.layers.Dense(10, activation='softmax'))


    model.compile(optimizer=tf.keras.optimizers.Adam(lr = 1e-3), loss="categorical_crossentropy", metrics=["accuracy"])


    x_train_, y_train_ = get_random_sample(x_train, y_train, number_of_samples=100)
    y_train_ = tf.keras.utils.to_categorical(y_train_)
    print(np.shape(y_train.shape))

    datagen = ImageDataGenerator(
        rotation_range = 15,
        zoom_range = 0.1,
        width_shift_range=2,
        height_shift_range=2,
        shear_range = 0.5,
        fill_mode = 'constant',
        cval = 0)
    print(model.summary())
#x_test_ = np.append(x_train[300:],x_test).reshape(x_train[300:].shape[0]+x_test.shape[0],28,28,1)
#y_test_ = np.append(y_train[300:],y_test).reshape(y_train[300:].shape[0]+y_test.shape[0],10)

#    csv_logger = CSVLogger('output/fashion_exacly_like_novatec__'+i+'.log')
#    history = model.fit(datagen.flow(x_train, tf.keras.utils.to_categorical(y_train), batch_size=20), validation_data=(x_test, y_test), epochs=125, steps_per_epoch = x_train_.shape[0]//20, callbacks=[csv_logger])
#    history = model.fit(datagen.flow(x_train, tf.keras.utils.to_categorical(y_train), batch_size=30),steps_per_epoch=2000,
#              validation_data=(x_test, y_test),
#              epochs=125, callbacks=[csv_logger],
#              shuffle=True)