import numpy as np
import pandas as pd 
import tensorflow as tf
import matplotlib.pyplot as plt

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()


# Let's visualize some of the images.
for i in range(4):
  plt.imshow(x_train[i])
  plt.show()

# Counting the examples for each class.
pd.DataFrame(y_train, columns=['labels'])['labels'].value_counts()

# Defining the model

# Defining the input.
inp = tf.keras.layers.Input(shape=(28,28,1))
# Rescaling the images.
resc = tf.keras.layers.Rescaling(1./255)(inp)
# Extracting the features.
conv = tf.keras.layers.Conv2D(filters=32, kernel_size=(4,4), activation='relu', input_shape=(28,28,1))(resc)
pool = tf.keras.layers.MaxPool2D(pool_size=(2,2))(conv)
flat = tf.keras.layers.Flatten()(pool)
# Interpreting the extracted features.
dense_1 = tf.keras.layers.Dense(units=100, activation='relu')(flat)
dense_2 = tf.keras.layers.Dense(units=50, activation='relu')(dense_1)
# Defining the output.
output = tf.keras.layers.Dense(units=10, activation='softmax')(dense_2)

model = tf.keras.Model(inputs=inp, outputs=output)

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



model.fit(x_train, y_train,
          epochs=10, 
          validation_data=(x_test, y_test))