import numpy as np
import tensorflow as tf
import tf_explain

INPUT_SHAPE = (28, 28, 1)
NUM_CLASSES = 10

AVAILABLE_DATASETS = {
    'mnist': tf.keras.datasets.mnist,
    'fashion_mnist': tf.keras.datasets.fashion_mnist,
}
DATASET_NAME = 'fashion_mnist'  # Choose between "mnist" and "fashion_mnist"

# Load dataset
dataset = AVAILABLE_DATASETS[DATASET_NAME]
(train_images, train_labels), (test_images, test_labels) = dataset.load_data()

# Convert from (28, 28) images to (28, 28, 1)
train_images = train_images[..., tf.newaxis]
test_images = test_images[..., tf.newaxis]

# One hot encore labels 0, 1, .., 9 to [0, 0, .., 1, 0, 0]
train_labels = tf.keras.utils.to_categorical(train_labels, num_classes=NUM_CLASSES)
test_labels = tf.keras.utils.to_categorical(test_labels, num_classes=NUM_CLASSES)

# Create model
img_input = tf.keras.Input(INPUT_SHAPE)

x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu')(img_input)
x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation='relu', name='target_layer')(x)
x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)

x = tf.keras.layers.Dropout(0.25)(x)
x = tf.keras.layers.Flatten()(x)

x = tf.keras.layers.Dense(128, activation='relu')(x)
x = tf.keras.layers.Dropout(0.5)(x)

x = tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')(x)

model = tf.keras.Model(img_input, x)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# Select a subset of the validation data to examine
# Here, we choose 5 elements with label "0" == [1, 0, 0, .., 0]
validation_class_zero = (np.array([
    el for el, label in zip(test_images, test_labels)
    if np.all(label == np.array([1] + [0] * 9))
][0:5]), None)
# Select a subset of the validation data to examine
# Here, we choose 5 elements with label "4" == [0, 0, 0, 0, 1, 0, 0, 0, 0, 0]
validation_class_fours = (np.array([
    el for el, label in zip(test_images, test_labels)
    if np.all(label == np.array([0] * 4 + [1] + [0] * 5))
][0:5]), None)

# Instantiate callbacks
# class_index value should match the validation_data selected above
callbacks = [
    tf_explain.callbacks.GradCAMCallback(validation_class_zero, 'target_layer', class_index=0),
    tf_explain.callbacks.GradCAMCallback(validation_class_fours, 'target_layer', class_index=4),
    tf_explain.callbacks.ActivationsVisualizationCallback(validation_class_zero, layers_name=['target_layer']),
    tf_explain.callbacks.SmoothGradCallback(validation_class_zero, class_index=0, num_samples=15, noise=1.),
    tf_explain.callbacks.IntegratedGradientsCallback(validation_class_zero, class_index=0, n_steps=10),
]

# Start training
model.fit(train_images, train_labels, epochs=5, callbacks=callbacks)

tensorflow.keras.Input tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2D tensorflow.keras.layers.MaxPool2D tensorflow.keras.Model tensorflow.keras.layers.Dropout numpy.array tensorflow.keras.layers.Flatten tensorflow.keras.utils.to_categorical
[(23, 'tensorflow.keras.utils.to_categorical', 'tf.keras.utils.to_categorical', (['train_labels'], {'num_classes': 'NUM_CLASSES'}), True, 'import tensorflow as tf\n'), (24, 'tensorflow.keras.utils.to_categorical', 'tf.keras.utils.to_categorical', (['test_labels'], {'num_classes': 'NUM_CLASSES'}), True, 'import tensorflow as tf\n'), (27, 'tensorflow.keras.Input', 'tf.keras.Input', (['INPUT_SHAPE'], {}), True, 'import tensorflow as tf\n'), (41, 'tensorflow.keras.Model', 'tf.keras.Model', (['img_input', 'x'], {}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', ([], {'filters': '(32)', 'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), True, 'import tensorflow as tf\n'), (30, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', ([], {'filters': '(64)', 'kernel_size': '(3, 3)', 'activation': '"""relu"""', 'name': '"""target_layer"""'}), True, 'import tensorflow as tf\n'), (31, 'tensorflow.keras.layers.MaxPool2D', 'tf.keras.layers.MaxPool2D', ([], {'pool_size': '(2, 2)'}), True, 'import tensorflow as tf\n'), (33, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['(0.25)'], {}), True, 'import tensorflow as tf\n'), (34, 'tensorflow.keras.layers.Flatten', 'tf.keras.layers.Flatten', ([], {}), True, 'import tensorflow as tf\n'), (36, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(128)'], {'activation': '"""relu"""'}), True, 'import tensorflow as tf\n'), (37, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['(0.5)'], {}), True, 'import tensorflow as tf\n'), (39, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['NUM_CLASSES'], {'activation': '"""softmax"""'}), True, 'import tensorflow as tf\n'), (60, 'tf_explain.callbacks.GradCAMCallback', 'tf_explain.callbacks.GradCAMCallback', (['validation_class_zero', '"""target_layer"""'], {'class_index': '(0)'}), False, 'import tf_explain\n'), (61, 'tf_explain.callbacks.GradCAMCallback', 'tf_explain.callbacks.GradCAMCallback', (['validation_class_fours', '"""target_layer"""'], {'class_index': '(4)'}), False, 'import tf_explain\n'), (62, 'tf_explain.callbacks.ActivationsVisualizationCallback', 'tf_explain.callbacks.ActivationsVisualizationCallback', (['validation_class_zero'], {'layers_name': "['target_layer']"}), False, 'import tf_explain\n'), (63, 'tf_explain.callbacks.SmoothGradCallback', 'tf_explain.callbacks.SmoothGradCallback', (['validation_class_zero'], {'class_index': '(0)', 'num_samples': '(15)', 'noise': '(1.0)'}), False, 'import tf_explain\n'), (64, 'tf_explain.callbacks.IntegratedGradientsCallback', 'tf_explain.callbacks.IntegratedGradientsCallback', (['validation_class_zero'], {'class_index': '(0)', 'n_steps': '(10)'}), False, 'import tf_explain\n'), (48, 'numpy.array', 'np.array', (['([1] + [0] * 9)'], {}), True, 'import numpy as np\n'), (54, 'numpy.array', 'np.array', (['([0] * 4 + [1] + [0] * 5)'], {}), True, 'import numpy as np\n')]