import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.layers import Flatten, Dense, Activation
from tensorflow.keras.models import Sequential

x_train = x_train.astype('float32')
x_test = x_test.astype('float32')

gray_scale = 255
x_train /= gray_scale
x_test /= gray_scale

print("Feature matrix:", x_train.shape)
print("Target matrix:", x_test.shape)
print("Feature matrix:", y_train.shape)
print("Target matrix:", y_test.shape)

fig, ax = plt.subplots(10, 10)
k = 0
for i in range(10):
    for j in range(10):
        ax[i][j].imshow(x_train[k].reshape(28, 28), aspect='auto')
        k += 1

model = Sequential([
    Flatten(input_shape=(28, 28)),
    Dense(10, activation='sigmoid') # single hidden layer
])


model.compile(optimizer='adam',loss='sparse_categorical_crossentropy',
              metrics='accuracy')

history = model.fit(x_train, y_train,
          epochs=50,
          batch_size=2000,
          validation_split=0.2)

results = model.evaluate(x_test,  y_test, verbose = 0)
print('For multilayer perceptron: test loss, test acc:', results)