# Import personnal set of tools
import sys
sys.path.append(r'../../Packages')
import SD.tools as SD_tools
from SD.tools import package_version as v

# Import other packages
from platform import python_version
print('Python version: {}'.format(python_version()))

import pandas as pd
v(pd)

import numpy as np
v(np)

import random

import os
import os.path
# os.environ["CUDA_VISIBLE_DEVICES"]="-1"  # Disable GPU. Instead use: "with tf.device('/cpu:0')" as needed

import matplotlib
import matplotlib.pyplot as plt
v(matplotlib)

import seaborn as sns

import sklearn
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score, confusion_matrix
v(sklearn)

import tensorflow as tf
from tensorflow.keras import backend as K
K.set_learning_phase(1)
v(tf)

from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import load_model
from tensorflow.keras.callbacks import ModelCheckpoint, CSVLogger
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import ReduceLROnPlateau, EarlyStopping

import pickle
import itertools
import keras # <---- only imported for SHAP, not used to run models. Use tf.keras instead. 

import shap
from shap.plots import colors 
v(shap)

# Disable warnings after notebook is completed
import warnings
warnings.filterwarnings('ignore')

# Define custom parameters

# In Tensorflow 2.0 val_acc was renamed to val_accuracy
acctxt = 'accuracy' if int(tf.__version__[0])>=2 else 'acc'

fsize=(15,5)          # Figure size for plots
n_examples=None       # limit examples for training, or None for all
nepochs=30            # Number of epochs for NN and CNN
target_accuracy=0.9962# Target for cross validation accuracy
lr_min=1e-8           # Minimum allowable learning rate
dropout_rate=0.30     # Dropout rate for all layers except final one
dropout_rate_ll=0.40  # Dropout rate for last layer
n_samples_max=500     # Maximum number of background samples to render SHAP values
train_logistic=1      # Flag to enable (re) training logistic models
train_nn=1            # Flag to enable (re) training neural network
train_cnn=1           # Flag to enable (re) training convoluted neural networks
advanced_vis=1        # Flag to allow advanced visualizations (CPU intensive)

# # # Preliminary runs parameters
# n_examples=293        # limit examples for training, or None for all
# nepochs=30            # Number of epochs for NN and CNN
# target_accuracy=0.92  # Target for cross validation accuracy
# train_logistic=0      # Flag to enable (re) training logistic models
# train_nn=1            # Flag to enable (re) training neural network
# train_cnn=1           # Flag to enable (re) training convoluted neural networks
# advanced_vis=0        # Flag to allow advanced visualizations (CPU intensive)

# # Preliminary extended runs parameters
# n_examples=1999       # limit examples for training, or None for all
# nepochs=30            # Number of epochs for NN and CNN
# target_accuracy=0.97  # Target for cross validation accuracy
# train_logistic=1      # Flag to enable (re) training logistic models
# train_nn=1            # Flag to enable (re) training neural network
# train_cnn=1           # Flag to enable (re) training convoluted neural networks
# advanced_vis=1        # Flag to allow advanced visualizations (CPU intensive)

(X_train, Y_train), (X_test, Y_test) = mnist.load_data()

# Provide basic information about the training and test set. 
print('The train set contains {} images of {}x{} pixels.\nThe test set contains {} images of {}x{} pixels.'.format(
    *(X_train.shape), *(X_test.shape)))
print('There are {} unique values in the set.'.format(len(np.unique(Y_train))))

# Print unique values from the train set. 
print('Unique values in the train set:')
unique = np.unique(Y_train)
print(unique)

# Plot samples from the train set. 
fig, axes = plt.subplots(10,10,figsize=(8,8))
for digit in unique:
    digit_images=X_train[Y_train==digit][20:]
    for k in range(10):
        axes[k,digit].imshow(digit_images[k], cmap='Greys')
        axes[k,digit].axis('off')
plt.suptitle('Samples from the dataset.', y=0.92);

# Plot the distribution of the unique classes in the train and test set. 
fig, axes = plt.subplots(1,2, figsize=fsize)
_, counts = np.unique(Y_train, return_counts=True)
axes[0].bar(unique, counts)
axes[0].set_xticks(unique)
axes[0].set_title('Distribution of unique classes in the training set.')
_, counts = np.unique(Y_test, return_counts=True)
axes[1].bar(unique, counts)
axes[1].set_xticks(unique)
axes[1].set_title('Distribution of unique classes in the test set.');

# Normalize and reshape data.

n=n_examples # limit number of examples, or None for all
(X_train, Y_train), (X_test, Y_test) = mnist.load_data()

# Reduce the number of examples for initial tests
X_train=X_train[:n]/255
X_test=X_test[:n]/255
Y_train=Y_train[:n]
Y_test=Y_test[:n]

# Backup and reshape sets
X_train0 = X_train
X_train = X_train.reshape(X_train.shape[0],-1)
X_test0 = X_test
X_test = X_test.reshape(X_test.shape[0],-1)

print('X_train set shape: {}x{}, Y_train set shape: {}(x1)'.format(*X_train.shape, *Y_train.shape))

# Dataframe containing the accuracies of the models.
df_res = pd.DataFrame(columns=['Classifier','Training accuracy','Testing accuracy'])

# Initialize and fit a linear regression model.
if train_logistic:
    clf=sklearn.linear_model.SGDClassifier(penalty='l1')
    clf.fit(X_train,Y_train)
    train_acc=accuracy_score(clf.predict(X_train), Y_train)
    test_acc=accuracy_score(clf.predict(X_test), Y_test)
    print('Linear logistic regression with Lasso regularization:')
    print('Train accuracy: {}, test accuracy: {}.'.format(round(train_acc,4),round(test_acc,4)))
    df_res=df_res.append({'Classifier': 'Linear - Lasso','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True);

# Initialize and fit a linear regression model.
if train_logistic:
    clf=sklearn.linear_model.SGDClassifier(penalty='l2')
    clf.fit(X_train,Y_train)
    train_acc=accuracy_score(clf.predict(X_train), Y_train)
    test_acc=accuracy_score(clf.predict(X_test), Y_test)
    print('Linear logistic regression with Ridge regularization:')
    print('Train accuracy: {}, test accuracy: {}.'.format(round(train_acc,4),round(test_acc,4)))
    df_res=df_res.append({'Classifier': 'Linear - Ridge','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True);

# Initialize and fit a linear regression model.
if train_logistic:
    clf=sklearn.linear_model.SGDClassifier(penalty='elasticnet')
    clf.fit(X_train,Y_train)
    train_acc=accuracy_score(clf.predict(X_train), Y_train)
    test_acc=accuracy_score(clf.predict(X_test), Y_test)
    print('Linear logistic regression with Elastic Net regularization:')
    print('Train accuracy: {}, test accuracy: {}.'.format(round(train_acc,4),round(test_acc,4)))
    df_res=df_res.append({'Classifier': 'Linear - ElasticNet','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True)

# Display the summary table.
df_res.round(4)

# Build and train the neural network.

model_file = 'nn_nodrop'

with tf.device('/cpu:0'):

    if train_nn or not os.path.isfile(model_file+'.h5'):

        nn_nodrop = tf.keras.Sequential([
            tf.keras.layers.Flatten(input_shape=(784,)),
            tf.keras.layers.Dense(256, activation='relu',),
            tf.keras.layers.Dense(128, activation='relu',),
            tf.keras.layers.Dense(64, activation='relu',),
            tf.keras.layers.Dense(32, activation='relu',),
            tf.keras.layers.Dense(10, activation='softmax',),
        ])
        print('Model defined.')

        # Compile the model
        nn_nodrop.compile(
            optimizer=tf.keras.optimizers.Adam(),
            loss='sparse_categorical_crossentropy',
            metrics=['acc'],)
        print('Model compiled.')

        # Train the model
        nn_nodrop_run=nn_nodrop.fit(X_train,
                      Y_train, 
                      epochs=nepochs,
                     validation_data=(X_test, Y_test),
                     )
        print('Model trained.')
        
        # Save the history on the local drive
        nn_nodrop_history = nn_nodrop_run.history
        with open(model_file+'.hist', 'wb') as file_pi:
            pickle.dump(nn_nodrop_history, file_pi)
            print('History saved.')

        # Save the model on the local drive
        nn_nodrop.save(model_file+'.h5')  # creates a HDF5 file 
        print('Model saved.')

    else:
        # Restore the model and its history from the local drive
        nn_nodrop=load_model(model_file+'.h5')
        nn_nodrop_history = pickle.load(open(model_file+'.hist', "rb"))
        print('Model and history reloaded.')

# Evaluate accuracy on the train and test set. 
train_acc=nn_nodrop.evaluate(X_train, Y_train)[1]
test_acc=nn_nodrop.evaluate(X_test, Y_test)[1]
print('\n\nTrain accuracy: {}% over {} epochs.'.format(round(train_acc*100,2),len(nn_nodrop_history['loss'])))
print('Test accuracy: {}% over {} epochs.'.format(round(test_acc*100,2),len(nn_nodrop_history['loss'])))
df_res=df_res.append({'Classifier': 'NN - no dopout','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True)

# Build and train the neural network.

model_file = 'nn_drop'

with tf.device('/cpu:0'):

    if train_nn or not os.path.isfile(model_file+'.h5'):

        nn_drop = tf.keras.Sequential([
        tf.keras.layers.Flatten(input_shape=(784,)),
        tf.keras.layers.Dense(256, activation='relu',),
        tf.keras.layers.Dropout(dropout_rate),
        tf.keras.layers.Dense(128, activation='relu',),
        tf.keras.layers.Dropout(dropout_rate),
        tf.keras.layers.Dense(64, activation='relu',),
        tf.keras.layers.Dropout(dropout_rate),
        tf.keras.layers.Dense(32, activation='relu',),
        tf.keras.layers.Dropout(dropout_rate),
        tf.keras.layers.Dense(10, activation='softmax',),
        ])
        print('Model defined.')

        # Compile the model
        nn_drop.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss='sparse_categorical_crossentropy',
        metrics=['acc'],
        )

        # Train the model
        nn_drop_run=nn_drop.fit(X_train,
                      Y_train, 
                      epochs=nepochs,
                     validation_data=(X_test, Y_test))
        print('Model trained.')
        
        # Save the history on the local drive
        nn_drop_history = nn_drop_run.history
        with open(model_file+'.hist', 'wb') as file_pi:
            pickle.dump(nn_drop_history, file_pi)
            print('History saved.')

        # Save the model on the local drive
        nn_drop.save(model_file+'.h5')  # creates a HDF5 file 
        print('Model saved.')

    else:
        # Restore the model and its history from the local drive
        nn_drop=load_model(model_file+'.h5')
        nn_drop_history = pickle.load(open(model_file+'.hist', "rb"))
        print('Model and history reloaded.')

# Evaluate accuracy on the train and test set. 
train_acc=nn_drop.evaluate(X_train, Y_train)[1]
test_acc=nn_drop.evaluate(X_test, Y_test)[1]
print('\n\nTrain accuracy: {}% over {} epochs.'.format(round(train_acc*100,2),len(nn_drop_history['loss'])))
print('Test accuracy: {}% over {} epochs.'.format(round(test_acc*100,2),len(nn_drop_history['loss'])))
df_res=df_res.append({'Classifier': 'NN - w/ dropout','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True)

# Plot the accuracies of the model versus the number of training epochs.

try:
    x=range(1, 1+nepochs)
    fig, ax = plt.subplots(figsize=fsize)
    ax = sns.lineplot(x, nn_nodrop_history[acctxt], color='b', label='Training accuracy - no dropout')
    ax = sns.lineplot(x, nn_nodrop_history['val_'+acctxt], color='g', label='Validation accuracy - no dropout')
    ax = sns.lineplot(x, nn_drop_history[acctxt], color='gray', label='Training accuracy - w/ dropout')
    ax = sns.lineplot(x, nn_drop_history['val_'+acctxt], color='black', label='Validation accuracy - w/ dropout')
    plt.title('Training and validation accuracy for neural networks')
    plt.ylabel('Accuracy')
    plt.xlabel('Epochs');
except:
    pass

# Compute predictions

X_test_pred=nn_drop.predict(X_test).argmax(axis=1) # Prediction of the model on X_test
y_pred=X_test_pred[X_test_pred!=Y_test] # Prediction of the model that are wrong
y_actual=Y_test[X_test_pred!=Y_test] # Ground truth corresponding to the wrong predictions
x=X_test0[X_test_pred!=Y_test] # X images corresponding to the wrong predictions

# Display miscategorized examples sorted by true class.

fig, axes = plt.subplots(10,10,figsize=(15,15))
fig.suptitle('Miscategorized examples sorted by true class.', fontsize=16, y=0.93);
for i in range(10):
    for j in range(10):
            y_actual=Y_test[np.logical_and(X_test_pred!=Y_test, Y_test==j)]
            y_pred=X_test_pred[np.logical_and(X_test_pred!=Y_test, Y_test==j)]
            x=X_test0[np.logical_and(X_test_pred!=Y_test, Y_test==j)]
            axes[i,j].axis('off')
            try:
                axes[i,j].imshow(x[i], cmap='Greys')
                axes[i,j].set_title('{} inst. of {}'.format(y_pred[i], y_actual[i]))
            except:
                pass

# Display miscategorized examples sorted by true class (x-axis) and prediction class (y-axis).

fig, axes = plt.subplots(10,10,figsize=(15,15))
fig.suptitle('Miscategorized examples sorted by true class (x-axis) and prediction class (y-axis).', fontsize=16, y=0.93);
for i in range(10):
    for j in range(10):
            axes[i,j].axis('off')
            if i==j: continue
            y_actual=Y_test[np.logical_and(X_test_pred==i, Y_test==j)]
            y_pred=X_test_pred[np.logical_and(X_test_pred==i, Y_test==j)]
            x=X_test0[np.logical_and(X_test_pred==i, Y_test==j)]
            try:
                axes[i,j].imshow(x[0], cmap='Greys')
                axes[i,j].set_title('{} inst. of {}'.format(i, j))
            except:
                pass

# Plot confusion matrices

try:
    # Plot confusion matrix
    fig, ax = plt.subplots(1,2,figsize=(16,6.5))
    shrink=1
    # plt.subplots_adjust(wspace=1)

    #############
    # Left plot #
    #############
    y_pred=nn_drop.predict(X_test).argmax(axis=1)
    y_actual=Y_test

    # Define confusion matrix
    cm = confusion_matrix(y_actual,y_pred)
    cm0 = cm
    cm = 100*cm.astype('float') / cm.sum(axis = 1)[:, np.newaxis]
    n_classes = len(unique)

    plt.sca(ax[0])
    plt.imshow(cm, cmap = 'YlOrRd')
    plt.title('Normalized confusion matrix (%)')
    plt.colorbar(shrink=shrink)
    tick_marks = np.arange(n_classes)
    plt.xticks(tick_marks, np.arange(n_classes))
    plt.yticks(tick_marks, np.arange(n_classes))
    thresh = cm.max() / 2.

    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        if i!=j:
            plt.text(j, i, format(cm[i, j], '.2f'),
                     horizontalalignment="center",
                     color="white" if cm[i, j] > thresh else "black")
        if i==j:
            plt.text(j, i, format(cm[i, j], '.2f'),
                     horizontalalignment="center",
                     color="white")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label');


    ##############
    # Right plot #
    ##############

    y_test_pred=nn_drop.predict(X_test).argmax(axis=1)
    y_pred=X_test_pred[y_test_pred!=Y_test]
    y_actual=Y_test[y_test_pred!=Y_test]
    x=X_test0[y_test_pred!=Y_test]

    # Define confusion matrix
    cm = confusion_matrix(y_actual,y_pred)
    cm = cm.astype('float')
    n_classes = len(unique)

    plt.sca(ax[1])
    thresh = cm.max() / 2.
    try:
        cm[range(10), range(10)] = np.nan
    except:
        pass
    plt.imshow(cm, cmap = 'YlOrRd')
    plt.title('Confusion matrix (number of images)')
    plt.colorbar(shrink=shrink)
    tick_marks = np.arange(n_classes)
    plt.xticks(tick_marks, np.arange(n_classes))
    plt.yticks(tick_marks, np.arange(n_classes))


    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        if i!=j:
            plt.text(j, i, format(cm[i, j], '.0f'),
                     horizontalalignment="center",
                     color="white" if cm[i, j] > thresh else "black")
        if i==j:
            plt.text(j, i, format(cm0[i, j], '.0f'),
                     horizontalalignment="center",
                     color="black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label');

except:
    pass

# Prepare the data (regularize and add axes as needed to fit the model input shape)

n=n_examples # limit examples, or None for all

(X_train, Y_train), (X_test, Y_test) = mnist.load_data()
X_train=X_train[:n]/255
X_test=X_test[:n]/255
Y_train=Y_train[:n]
Y_test=Y_test[:n]
X_train = X_train[:,:,:,np.newaxis]
X_test = X_test[:,:,:,np.newaxis]

# Define, compile and train the model.

model_file = 'cnn'

if train_cnn or not os.path.isfile(model_file+'.h5'):
        
    # Define the CNN model
    cnn=tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding = 'same', input_shape=(28,28,1)), # input shape: 28x28x1
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding = 'same'), # input shape: 26x26x32
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(32, (5,5), strides=2, activation='relu', padding='same'), # input shape: 24x24x32
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dropout(rate=dropout_rate),
    
    tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding = 'same'), # input shape: 12x12x32
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding = 'same'), # input shape: 10x10x64
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(64, (5,5), strides=2, activation='relu', padding='same'), # input shape: 8x8x64
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dropout(rate=dropout_rate),
    
    tf.keras.layers.Flatten(), # input shape: 4x4x64
     
    tf.keras.layers.Dense(128, activation='relu'), # input shape: 1024
    # tf.keras.layers.BatchNormalization(),  # <----- this layer was causing incompatibilities with SHAP
    tf.keras.layers.Dropout(rate=dropout_rate_ll),
        
    tf.keras.layers.Dense(10, activation='softmax') # input shape: 128
    ])


    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, amsgrad=False)
    
    # Compile the CNN model
    cnn.compile(
        optimizer=optimizer,
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    #Delete the training log if it exists. The training log is used to keep track of the losses and accuracies
    #over batch trainings.
    filepath=model_file+'.h5'
    try:
        os.remove(filepath)
        os.remove('cnn_training.log')
    except:
        pass

    
    checkpoint=ModelCheckpoint(filepath, monitor='loss', verbose=0, save_best_only=False, mode='auto')

    csv_logger = CSVLogger('cnn_training.log', append=True)
    
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2,
                              patience=4, min_lr=lr_min, verbose=1)

    callbacks=[checkpoint, csv_logger, reduce_lr]
    
    cnn_run=cnn.fit(
        X_train,
        Y_train,
        batch_size=32,
        epochs=nepochs,
        validation_data=(X_test, Y_test),
        callbacks=callbacks,
    )

else:
    # Restore the model and its history from the local drive
    cnn=load_model(model_file+'.h5')
    print('Model and history reloaded.')
    
with open('cnn_training.log', 'r') as f:
        lines=f.readlines()
        print('\n\nAccuracy: {}% over {} epochs.'.format(round(cnn.evaluate(X_test, Y_test)[1]*100,2),len(lines)-1))

# Evaluate accuracy on the train and test set. 

train_acc=cnn.evaluate(X_train, Y_train)[1]
test_acc=cnn.evaluate(X_test, Y_test)[1]
with open('cnn_training.log', 'r') as f:
    lines=f.readlines()
    print('\n\nTrain accuracy: {}% over {} epochs.'.format(round(train_acc*100,2),len(lines)-1))
    print('Test accuracy: {}% over {} epochs.'.format(round(test_acc*100,2),len(lines)-1))
df_res=df_res.append({'Classifier': 'CNN - no data augment','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True)

# Prepare the data (regularize and add axes as needed to fit the model input shape)

n=n_examples # limit examples, or None for all

(X_train, Y_train), (X_test, Y_test) = mnist.load_data()
X_train=X_train[:n]/255
X_test=X_test[:n]/255
Y_train=Y_train[:n]
Y_test=Y_test[:n]
X_train = X_train[:,:,:,np.newaxis]
X_test = X_test[:,:,:,np.newaxis]

# Define, compile and train the model.

model_file = 'cnn_augment'
filepath=model_file+'.h5'

if train_cnn or not os.path.isfile(model_file+'.h5'):
        
        # Data augmentation

    datagen=ImageDataGenerator(
        zoom_range=0.1,
        rotation_range=10,
        width_shift_range=0.1,
        height_shift_range=0.1
        )

    train_set_augment=datagen.flow(X_train, Y_train, batch_size=32)
    test_set_augment=datagen.flow(X_test, Y_test, batch_size=32)
    
    cnn_augment=tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding = 'same', input_shape=(28,28,1)), # input shape: 28x28x1
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding = 'same'), # input shape: 26x26x32
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(32, (5,5), strides=2, activation='relu', padding='same'), # input shape: 24x24x32
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dropout(rate=dropout_rate),
    
    tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding = 'same'), # input shape: 12x12x32
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding = 'same'), # input shape: 10x10x64
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Conv2D(64, (5,5), strides=2, activation='relu', padding='same'), # input shape: 8x8x64
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dropout(rate=dropout_rate),
    
    tf.keras.layers.Flatten(), # input shape: 4x4x64
     
    tf.keras.layers.Dense(128, activation='relu'), # input shape: 1024
    # tf.keras.layers.BatchNormalization(),  # <----- this layer was causing incompatibilities with SHAP
    tf.keras.layers.Dropout(rate=dropout_rate_ll),
        
    tf.keras.layers.Dense(10, activation='softmax') # input shape: 128
    ])
    
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, amsgrad=False)
    
    # Compile the CNN model
    cnn_augment.compile(
        optimizer=optimizer,
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
        )

    #Delete the training log if it exists. The training log is used to keep track of the losses and accuracies
    #over batch trainings.

    
    try:
        os.remove(filepath)
    except:
        pass
    
    try:
        os.remove('cnn_augment_training.log')
    except:
        pass
    
    checkpoint=ModelCheckpoint(filepath, monitor='loss', verbose=0, save_best_only=False, mode='auto')

    csv_logger = CSVLogger('cnn_augment_training.log', append=True)

    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2,
                              patience=4, min_lr=lr_min, verbose=1)

    callbacks=[checkpoint, csv_logger, reduce_lr]
    
    cnn_augment_run=cnn_augment.fit_generator(
        train_set_augment,
        epochs=nepochs,
        verbose=1,
        callbacks=callbacks,
        validation_data=(X_test, Y_test),
    )

else:
    # Restore the model and its history from the local drive
    cnn_augment=load_model(model_file+'.h5')
    checkpoint=ModelCheckpoint(filepath, monitor='loss', verbose=0, save_best_only=False, mode='auto')
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2,
                              patience=4, min_lr=lr_min, verbose=1)
    csv_logger = CSVLogger('cnn_augment_training.log', append=True)
    callbacks=[checkpoint, csv_logger, reduce_lr]
    print('Model and history reloaded.')

# Evaluate accuracy on the train and test set. 

train_acc=cnn_augment.evaluate(X_train, Y_train)[1]
test_acc=cnn_augment.evaluate(X_test, Y_test)[1]
with open('cnn_augment_training.log', 'r') as f:
    ep=len(f.readlines())-1
    print('\n\nAccuracy: {}% over {} epochs.'.format(round(train_acc*100,2),ep))
    print('Test accuracy: {}% over {} epochs.'.format(round(test_acc*100,2),ep))
df_res=df_res.append({'Classifier': 'CNN - w/ data augment','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True)

# Plot the accuracies of the model versus the number of training epochs.

x=range(1, 1+ep)
fig, ax = plt.subplots(figsize=fsize)
ax = sns.lineplot(x, cnn_run.history[acctxt], color='b', label='Training accuracy - no data augm')
ax = sns.lineplot(x, cnn_run.history['val_'+acctxt], color='g', label='Validation accuracy - no data augm')
ax = sns.lineplot(x, cnn_augment_run.history[acctxt], color='gray', label='Training accuracy - w/ data augment')
ax = sns.lineplot(x, cnn_augment_run.history['val_'+acctxt], color='black', label='Validation accuracy - w/ data augment')

# Compute predictions.

X_test_pred=cnn_augment.predict(X_test).argmax(axis=1)
y_pred=X_test_pred[X_test_pred!=Y_test]
y_actual=Y_test[X_test_pred!=Y_test]
x=X_test0[X_test_pred!=Y_test]

# Display miscategorized examples sorted by true class.

fig, axes = plt.subplots(10,10,figsize=(15,15))
fig.suptitle('Miscategorized examples sorted by true class.', fontsize=16, y=0.93);
for i in range(10):
    for j in range(10):
            y_actual=Y_test[np.logical_and(X_test_pred!=Y_test, Y_test==j)]
            y_pred=X_test_pred[np.logical_and(X_test_pred!=Y_test, Y_test==j)]
            x=X_test0[np.logical_and(X_test_pred!=Y_test, Y_test==j)]
            axes[i,j].axis('off')
            try:
                axes[i,j].imshow(x[i], cmap='Greys')
                axes[i,j].set_title('{} inst. of {}'.format(y_pred[i], y_actual[i]))
            except:
                pass

# Display miscategorized examples sorted by true class (x-axis) and prediction class (y-axis).

fig, axes = plt.subplots(10,10,figsize=(15,15))
fig.suptitle('Miscategorized examples sorted by true class (x-axis) and prediction class (y-axis).', fontsize=16, y=0.93);
for i in range(10):
    for j in range(10):
            axes[i,j].axis('off')
            if i==j: continue
            y_actual=Y_test[np.logical_and(X_test_pred==i, Y_test==j)]
            y_pred=X_test_pred[np.logical_and(X_test_pred==i, Y_test==j)]
            x=X_test0[np.logical_and(X_test_pred==i, Y_test==j)]
            try:
                axes[i,j].imshow(x[0], cmap='Greys')
                axes[i,j].set_title('{} inst. of {}'.format(i, j))
            except:
                pass

# Plot confusion matrices

fig, ax = plt.subplots(1,2,figsize=(16,6.5))
shrink=1
# plt.subplots_adjust(wspace=1)

#############
# Left plot #
#############
y_pred=cnn_augment.predict(X_test).argmax(axis=1)
y_actual=Y_test

# Define confusion matrix
cm = confusion_matrix(y_actual,y_pred)
cm0 = cm
cm = 100*cm.astype('float') / cm.sum(axis = 1)[:, np.newaxis]
n_classes = len(unique)

plt.sca(ax[0])
plt.imshow(cm, cmap = 'YlOrRd')
plt.title('Normalized confusion matrix (%)')
plt.colorbar(shrink=shrink)
tick_marks = np.arange(n_classes)
plt.xticks(tick_marks, np.arange(n_classes))
plt.yticks(tick_marks, np.arange(n_classes))
thresh = cm.max() / 2.

for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
    if i!=j:
        plt.text(j, i, format(cm[i, j], '.2f'),
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")
    if i==j:
        plt.text(j, i, format(cm[i, j], '.2f'),
                 horizontalalignment="center",
                 color="white")

plt.tight_layout()
plt.ylabel('True label')
plt.xlabel('Predicted label');


##############
# Right plot #
##############

y_test_pred=cnn_augment.predict(X_test).argmax(axis=1)
y_pred=X_test_pred[y_test_pred!=Y_test]
y_actual=Y_test[y_test_pred!=Y_test]
x=X_test0[y_test_pred!=Y_test]

# Define confusion matrix
cm = confusion_matrix(y_actual,y_pred)
cm = cm.astype('float')
n_classes = len(unique)

plt.sca(ax[1])
thresh = cm.max() / 2.

try:
    cm[range(10), range(10)] = np.nan
except:
    pass

plt.imshow(cm, cmap = 'YlOrRd')
plt.title('Confusion matrix (number of images)')
plt.colorbar(shrink=shrink)
tick_marks = np.arange(n_classes)
plt.xticks(tick_marks, np.arange(n_classes))
plt.yticks(tick_marks, np.arange(n_classes))


for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
    if i!=j:
        plt.text(j, i, format(cm[i, j], '.0f'),
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")
    if i==j:
        plt.text(j, i, format(cm0[i, j], '.0f'),
                 horizontalalignment="center",
                 color="black")

plt.tight_layout()
plt.ylabel('True label')
plt.xlabel('Predicted label');

# Display SHAP values.

try:
    if advanced_vis:

        with tf.device('/cpu:0'):

            # The following code is taken from https://github.com/slundberg/shap/blob/b606ab179d5b70ec6bd3e5acbaaed4c9bd65a14e/shap/plots/image.py
            # and slightly edited for better appearance of the colorbar. 

            try:
                nbackgroundsamples=min(n,n_samples_max) # <--- run time is proportionnal to that number!
            except:
                nbackgroundsamples=n_samples_max
            ntodisplay = 10

            # tf.compat.v1.Session()

            # select a set of background examples to take an expectation over
            rand=np.random.choice(X_train.shape[0], nbackgroundsamples, replace=False)
            background = X_train[rand]
            # print(background.shape)
            # print(Y_train[rand])

            # explain predictions of the model on three images
            e = shap.DeepExplainer(cnn_augment, background)

            # Select only example correctly classified
            bool1=Y_test==cnn_augment.predict(X_test).argmax(axis=1)
            X_test_correct=X_test[bool1]
            Y_test_correct=Y_test[bool1]
            X_select=[]
            for number in range(10):
                try:
                    X_select.append(X_test_correct[Y_test_correct==int(number)][0])
                except:
                    pass
            X_select=np.array(X_select)

            cnn_augment.summary()

            shap_values = e.shap_values(X_select)

            # # plot the feature attributions

            SD_tools.plot_shap_values(shap_values, -X_select)
            
except:
    print('Problem with SHAP.')

# Display SHAP values.

try:
    if advanced_vis:
        with tf.device('/cpu:0'):

            # tf.compat.v1.disable_v2_behavior()

            try:
                nbackgroundsamples=min(n,n_samples_max) # <--- run time is proportionnal to that number!
            except:
                nbackgroundsamples=n_samples_max

            ntodisplay = 10

            # select a set of background examples to take an expectation over
            rand=np.random.choice(X_train.shape[0], nbackgroundsamples, replace=False)
            background = X_train[rand]
            # print(Y_train[rand])

            # explain predictions of the model on three images
            e = shap.DeepExplainer(cnn_augment, background)

            # Select only example WRONGLYclassified
            bool1=Y_test!=cnn_augment.predict(X_test).argmax(axis=1)
            X_test_incorrect=X_test[bool1]
            Y_test_incorrect_predicted=cnn_augment.predict(X_test).argmax(axis=1)[bool1]
            Y_test_ground=Y_test[bool1]
            X_select=[]
            Y_select=[]
            for number in range(10):
                try:
                    X_select.append(X_test_incorrect[Y_test_ground==int(number)][0])
                    Y_select.append(Y_test_incorrect_predicted[Y_test_ground==int(number)][0])
                except:
                    pass
            X_select=np.array(X_select)
            print('Ground truth labels:')
            print(Y_select)

            shap_values = e.shap_values(X_select)

            # plot the feature attributions

            SD_tools.plot_shap_values(shap_values, -X_select)
        
        
except:
    print('Problem with SHAP.')

# Visualizating filters of the first layer
# https://blog.keras.io/how-convolutional-neural-networks-see-the-world.html

if advanced_vis:
    try:
        model=cnn_augment
        l2v = 0 # Layer to visualize
        SD_tools.visualize_filters(model = model, img = np.array(X_train[10]).reshape((1, 28, 28, 1)).astype(np.float64), 
                              layer_name = model.layers[l2v].name, print_summary=1, h_pad=0.05)
    except:
        pass

# Display the summary table.

print('The training and test accuracies for all models in thus notebook are shown below:')
df_res.round(5)

# Write conclusion.

max_test_accuracy=df_res['Testing accuracy'].max()
max_test_accuracy_row=df_res['Testing accuracy'].idxmax()
print('The model with the highest testing accuracy is "{}" with {:.2f}%'.format(df_res.iloc[max_test_accuracy_row,0], 100*max_test_accuracy))

# Prints test accuracy goal.

print('Target test accuracy goal: {:.2%}'.format(target_accuracy))

# Define an early stopping callback based in a target test set accuracy.

class EarlyStoppingByAccVal(tf.keras.callbacks.Callback):
    def __init__(self, monitor='val_acc', value=target_accuracy, verbose=1):
        super(tf.keras.callbacks.Callback, self).__init__()
        self.monitor = monitor
        self.value = value
        self.verbose = verbose

    def on_epoch_end(self, epoch, logs={}):
        current = logs.get(self.monitor)
        if current is None:
            warnings.warn("Early stopping requires %s available!" % self.monitor, RuntimeWarning)

        if current >= self.value:
            if self.verbose > 0:
                print("Epoch %05d: early stopping THR" % epoch)
            self.model.stop_training = True

# Perform additional training to attain target accuracy on test set

additional_training_epochs=1000

reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2,
                              patience=4, min_lr=1e-20, verbose=1)

callbacks=[checkpoint, csv_logger, reduce_lr]
    
    
callbacks_with_target=callbacks.copy()
callbacks_with_target.append(EarlyStoppingByAccVal())

if additional_training_epochs!=0:

    cnn_augment.fit_generator(
        train_set_augment,
        epochs=additional_training_epochs,
#         verbose=1,
        callbacks=callbacks_with_target,
#         validation_data=test_set_augment,
        validation_data=(X_test, Y_test),
    )

    with open('cnn_augment_training.log', 'r') as f:
        ep=len(f.readlines())-1
        print('\n\nAccuracy: {}% over {} epochs.'.format(round(cnn_augment.evaluate(X_test, Y_test)[1]*100,2),ep))

# Evaluate accuracy on the train and test set. 
from IPython.display import clear_output

i=0
test_acc=0

while test_acc < round(target_accuracy,3):
    i+=1
    print(i)
    if i%5==0:
        clear_output(wait=True)
    test_acc=cnn_augment.evaluate(X_test, Y_test)[1]
train_acc=cnn_augment.evaluate(X_train, Y_train)[1]
print('{} evaluations until target met (see explanations below)'.format(str(i)))
        
with open('cnn_augment_training.log', 'r') as f:
        ep=len(f.readlines())-1
        print('\n\nAccuracy: {}% over {} epochs.'.format(round(train_acc*100,2),ep))
        print('Test accuracy: {}% over {} epochs.'.format(round(test_acc*100,2),ep))

df_res=df_res.append({'Classifier': 'CNN - w/ DA and target','Training accuracy':train_acc,'Testing accuracy':test_acc}, ignore_index=True)

# Display the summary table.

print('The training and test accuracies for all models in this notebook are shown below:')
df_res.round(5)

# Write conclusion.

max_test_accuracy=df_res['Testing accuracy'].max()
max_test_accuracy_row=df_res['Testing accuracy'].idxmax()
print('The model with the highest testing accuracy is "{}" with {:.2f}%'.format(df_res.iloc[max_test_accuracy_row,0], 100*max_test_accuracy))