"""
 Trains a simple dense net on the MNIST dataset
 https://www.tensorflow.org/datasets/keras_example

 Needs modules tensorflow and tensorflow_datasets
"""

# silence tensorflow import verbose messages:
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt
import numpy as np




# Input pipeline
# --------------
(ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['train', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)

def plotit(ds):
    for image,label in ds:  
        plt.imshow(image, cmap='Greys')
        plt.title(label.numpy())
        plt.draw()
        plt.pause(0.1)


# Plot a few
# ----------
if True:
    print('plotting 10 samples from dataset')
    ds = ds_train.take(10) # .cache().repeat()
    plotit(ds)



# Training pipeline
# -----------------
def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label

ds_train = ds_train.map(
    normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
ds_train = ds_train.cache()
ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
ds_train = ds_train.batch(128)
ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)


# Evaluation pipeline
# -------------------

ds_test = ds_test.map(
    normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
ds_test = ds_test.batch(128)
ds_test = ds_test.cache()
ds_test = ds_test.prefetch(tf.data.experimental.AUTOTUNE)

# Define the model
# ----------------
model = tf.keras.models.Sequential(
    [
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(128,activation='relu'),
        tf.keras.layers.Dense(10)
    ])

model.summary()

# Compile the model
# -----------------
model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
)

# Train the model
# ---------------
model.fit(
    ds_train,
    epochs=6,
    validation_data=ds_test,
)

# what was badly identified
# -------------------------
ds_try = ds_test.take(5) # note: batches of 128 

pred_values = model.predict(ds_try) # returns numpy array
# convert output to categories 0...9
pred_values = np.argmax(pred_values, axis=1)

for im,lab in ds_try:
    true_values = lab.numpy()
    images = im

print(true_values.shape)
print(images.shape)

for image,value,pred in zip(images,true_values,pred_values):
    if value != pred:
        plt.imshow(image, cmap='Greys')
        plt.title(f'true value {value} model prediction {pred}')
        plt.draw()
        plt.pause(4.)