# Partially based on https://github.com/tensorflow/docs/blob/master/site/en/tutorials/keras/classification.ipynb

import tensorflow as tf

# Import mlpug for Tensorflow backend
import mlpug.tensorflow as mlp

# ################ SETTINGS ################
batch_size = 32
num_epochs = 10

progress_log_period = 500

tf.random.set_seed(0)  # For reproducibility
mlp.logging.use_fancy_colors()
# ##########################################

# ############## PROCESS DATA ##############
# Load data
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

# Scale
train_images = train_images / 255.0
test_images = test_images / 255.0
# ##########################################

# ########## SETUP BATCH DATASETS ##########
training_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).batch(batch_size)

# Using the test set as a validation set, just for demonstration purposes
validation_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(batch_size)
# ##########################################

# ############ BUILD THE MODEL #############
classifier = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])


# MLPug needs a TrainModel that outputs the loss
class TrainModel(tf.keras.Model):
    def __init__(self, classifier):
        super(TrainModel, self).__init__()

        self.classifier = classifier
        self.loss_func = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

    def call(self, batch_data, evaluate_settings, inference_mode=None):
        images, true_labels = batch_data

        logits = self.classifier(images)
        return self.loss_func(true_labels, logits)


train_model = TrainModel(classifier)
# ##########################################

# ############ SETUP OPTIMIZER #############
optimizer = tf.keras.optimizers.Adam()
# ##########################################

# ############# SETUP TRAINING ##############
trainer = mlp.trainers.DefaultTrainer(optimizers=optimizer,
                                      model_components=classifier,
                                      batch_data_signature=(tf.TensorSpec(shape=(None, 28, 28), dtype=tf.float64),
                                                            tf.TensorSpec(shape=(None,), dtype=tf.uint8),))

# At minimum you want to log the loss in the training progress
# By default the batch loss and the moving average of the loss are calculated and logged
loss_evaluator = mlp.evaluation.MetricEvaluator(trainer=trainer)
callbacks = [
    mlp.callbacks.TrainingMetricsLogger(metric_evaluator=loss_evaluator),
    # Calculate validation loss only once per epoch over the whole dataset
    mlp.callbacks.TestMetricsLogger(validation_dataset,
                                    'validation',
                                    metric_evaluator=loss_evaluator,
                                    batch_level=False),
    mlp.callbacks.LogProgress(log_period=progress_log_period, set_names=['training', 'validation']),
]

manager = mlp.trainers.TrainingManager(trainer,
                                       training_dataset,
                                       num_epochs=num_epochs,
                                       num_batches_per_epoch=int(training_dataset.cardinality()),
                                       callbacks=callbacks)

trainer.set_training_model(train_model)
# ##########################################

if __name__ == '__main__':
    # ################# START! #################
    manager.start_training()
    # ##########################################

    # ######### USE THE TRAINED MODEL ##########
    print("\nUsing the classifier ...")
    image = test_images[0]
    real_label = test_labels[0]

    logits = classifier(tf.expand_dims(image, 0))
    probabilities = tf.nn.softmax(logits)

    predicted_label = tf.math.argmax(probabilities, axis=-1)

    print(f"real label = {real_label}, predicted label = {tf.squeeze(predicted_label)}\n")
