# Load data
(X_train, y_train), (X_test, y_test) = load_my_data()
# Print an example
#print_image(X_train, 42)

# Normalize the inputs
X_train = X_train / 255.
X_test = X_test / 255.

# Define the optimizers to test
my_optimizers = {"Mini-batch GD":tf.keras.optimizers.SGD(learning_rate = 0.001, momentum = 0.0),
                 "Momentum GD":tf.keras.optimizers.SGD(learning_rate = 0.001, momentum = 0.9),
                 "RMS Prop":tf.keras.optimizers.RMSprop(learning_rate = 0.001, rho = 0.9),
                 "Adam":tf.keras.optimizers.Adam(learning_rate = 0.001, beta_1 = 0.9, beta_2 = 0.999)
    }

histories = {}
for optimizer_name, optimizer in my_optimizers.items():
    # Define a neural network

    my_network = tf.keras.models.Sequential([
                tf.keras.layers.Flatten(input_shape=(28, 28)),                    
                tf.keras.layers.Dense(128, activation='relu'),
                tf.keras.layers.Dense(64, activation='relu'),
                tf.keras.layers.Dense(10, activation='softmax')
                ])
    # Compile the model
    my_network.compile(optimizer=optimizer,
                       loss='sparse_categorical_crossentropy', # since labels are more than 2 and not one-hot-encoded
                       metrics=['accuracy'])

    # Train the model
    print('Training the model with optimizer {}'.format(optimizer_name))
    histories[optimizer_name] = my_network.fit(X_train, y_train, epochs=50, validation_split=0.1, verbose=1)

# Plot learning curves
for optimizer_name, history in histories.items():
    loss = history.history['loss']
    epochs = range(1,len(loss)+1)
    plt.plot(epochs, loss, label=optimizer_name)
    plt.legend(loc="upper right")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")