# -*- coding: utf-8 -*-
"""MNIST_fashionMNIST_training.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Ifw4dtUitUh87iiFkisFd0YGL-epdKDz
"""

datasetName = "MNIST"
# datasetName = "FashionMNIST"

from tensorflow.keras import layers, models

def getModel():
    model =  models.Sequential([
                layers.Conv2D(32, (5, 5), activation='relu', input_shape=(28, 28, 1)),
                layers.MaxPooling2D((2, 2) , strides = (2,2)),
                layers.Conv2D(32, (5, 5), activation='relu'),
                layers.MaxPooling2D((2, 2) , strides = (2,2)),
                layers.Flatten(),
                layers.Dense(100, activation='relu'),
                layers.Dense(10, activation='softmax')])
    model.compile(optimizer='sgd', loss='categorical_crossentropy', metrics=['accuracy'])
    return model

# print(getModel().summary())

import tensorflow as tf

if (datasetName == "MNIST"):
    (train_data, train_labels), (test_data, test_labels) = tf.keras.datasets.mnist.load_data()
if (datasetName == "FashionMNIST"):
    (train_data, train_labels), (test_data, test_labels) = tf.keras.datasets.fashion_mnist.load_data()


#reshaping
test_data = test_data.reshape((test_data.shape[0], 28, 28, 1))
train_data = train_data.reshape((train_data.shape[0], 28, 28, 1))
test_labels = test_labels.reshape(test_labels.shape[0])
train_labels = train_labels.reshape(train_labels.shape[0])

# convert from integers to floats
train_data = train_data.astype('float32')
test_data = test_data.astype('float32')

# normalize to range 0-1
train_data = train_data / 255.0
test_data = test_data / 255.0

import numpy as np

noisePerc = 20 

noisy_labels = np.load(str(noisePerc) + "_NoisyLabels_" + datasetName + ".npy")   #load noisy labels
noisy_lvl = np.load(str(noisePerc) + "_NoiseLevels_" + datasetName +".npy")  # load noise levels
filteredPrediction = np.load(str(noisePerc) + "_NoiseLevelPrediction_" + datasetName +".npy")  #load filtered predictions

grn_truth = np.array(noisy_labels == train_labels, dtype=int)
# print("Number of mislabelled: ", len(grn_truth) - sum(grn_truth), "out of", len(grn_truth))

from random import sample


if (datasetName == "MNIST"):
    sampleSize = int((5/100 * sum(filteredPrediction)))
else : 
    sampleSize = int(sum(filteredPrediction))

idx = sample([i for i in range(int(sum(filteredPrediction)))], sampleSize)

train_y = tf.keras.utils.to_categorical(noisy_labels, 10)
test_y = tf.keras.utils.to_categorical(test_labels, 10)

idxFilt = [i for i in range(len(grn_truth)) if filteredPrediction[i] == 1]

model = getModel()
earlyStopCallback = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
splitID = int(0.8 * sampleSize)

hist = model.fit(train_data[idxFilt][idx][:splitID], train_y[idxFilt][idx][:splitID],
                 epochs = 100, 
                 validation_data = (train_data[idxFilt][idx][splitID:], train_y[idxFilt][idx][splitID:]),
                 callbacks=[earlyStopCallback], shuffle=True, batch_size=32, verbose=1)

model.evaluate(test_data, test_y)