"""
Credits: This code is adapted from code originally posted 
by François Chollet at URL 
https://keras.io/examples/vision/visualizing_what_convnets_learn/ . 
"""

#%%

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.applications.resnet_v2 import preprocess_input, decode_predictions
from tensorflow.keras.preprocessing import image
from IPython.display import Image, display

#%%

dataset = keras.datasets.mnist

(x_train, train_labels), (x_test, test_labels) = dataset.load_data()
number_of_classes = np.max([np.max(train_labels), np.max(test_labels)]) + 1

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255

# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")
input_shape = x_train[0].shape

"""
#%%

model = keras.Sequential(    [
        keras.Input(shape=input_shape),
        keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        keras.layers.MaxPooling2D(pool_size=(2, 2)),
        keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        keras.layers.MaxPooling2D(pool_size=(2, 2)),
        keras.layers.Flatten(),
        keras.layers.Dropout(0.5),
        keras.layers.Dense(number_of_classes, activation="softmax"),
    ])

model.summary()

batch_size = 128
epochs = 15

model.compile(loss=keras.losses.SparseCategoricalCrossentropy(), optimizer="adam", metrics=["accuracy"])


model.fit(x_train, train_labels, epochs=epochs, batch_size=batch_size)
test_loss, test_acc = model.evaluate(x_test,  test_labels, verbose=2)
print('\nTest accuracy: %.2f%%' % (test_acc * 100))

#%%

model.save("mnist_cnn_1.keras")
"""

#%%

filename = "mnist_cnn_1.keras"
model = keras.models.load_model(filename)
model.summary()

# The dimensions of our input image
(img_height, img_width) = x_train.shape[1:3]

#%%

layer_name = "conv2d"
#layer_name = "conv2d_1"

# Set up a model that returns the activation values for our target layer
layer = model.get_layer(name=layer_name)
truncated_model = keras.Model(inputs=model.inputs, outputs=layer.output)

#%%

#@tf.function
def gradient_ascent_step(img, filter_index, learning_rate):
    with tf.GradientTape() as tape:
        tape.watch(img)
        activation = truncated_model(img)
        filter_activation = activation[:, 2:-2, 2:-2, filter_index]
        gain = tf.reduce_mean(filter_activation)

    # Compute gradients.
    grads = tape.gradient(gain, img)
    # Normalize gradients.
    grads = tf.math.l2_normalize(grads)
    img += learning_rate * grads
    return gain, img


def initialize_image():
    # We start from a gray image with some random noise
    img = tf.random.uniform((1, img_width, img_height, 1))
    # ResNet50V2 expects inputs in the range [-1, +1].
    # Here we scale our random inputs to [-0.125, +0.125]
    return img * 0.25


def visualize_filter(filter_index):
    # We run gradient ascent for 30 steps
    iterations = 30
    learning_rate = 10.0
    img = initialize_image()
    for iteration in range(iterations):
        gain, img = gradient_ascent_step(img, filter_index, learning_rate)

    # Decode the resulting input image
    img = deprocess_image(img[0].numpy())
    return gain, img


def deprocess_image(img):
    # Normalize array: center on 0., ensure variance is 0.15
    img -= img.mean()
    img /= img.std() + 1e-5
    img *= 0.15

    # Center crop
    #img = img[25:-25, 25:-25, :]

    # Clip to [0, 1]
    img += 0.5
    img = np.clip(img, 0, 1)

    # Convert to RGB array
    img *= 255
    img = np.clip(img, 0, 255).astype("uint8")
    return img


#%%

filter_index = -1

#%%

filter_index = filter_index + 1
gain, img = visualize_filter(filter_index)
filename = "0_%d.png" % (filter_index)
keras.preprocessing.image.save_img(filename, img)
display(Image(filename))
print("filter_index = ", filter_index)


#%%

# Compute image inputs that maximize per-filter activations
# for the first 64 filters of our target layer
all_imgs = []
for filter_index in range(64):
    print("Processing filter %d" % (filter_index,))
    gain, img = visualize_filter(filter_index)
    all_imgs.append(img)

# Build a black picture with enough space for
# our 8 x 8 filters of size 128 x 128, with a 5px margin in between
margin = 5
n = 8
cropped_width = img_width - 25 * 2
cropped_height = img_height - 25 * 2
width = n * cropped_width + (n - 1) * margin
height = n * cropped_height + (n - 1) * margin
stitched_filters = np.zeros((width, height, 3))

# Fill the picture with our saved filters
for i in range(n):
    for j in range(n):
        img = all_imgs[i * n + j]
        stitched_filters[
            (cropped_width + margin) * i : (cropped_width + margin) * i + cropped_width,
            (cropped_height + margin) * j : (cropped_height + margin) * j
            + cropped_height,
            :,
        ] = img
keras.preprocessing.image.save_img("stiched_filters.png", stitched_filters)

from IPython.display import Image, display

display(Image("stiched_filters.png"))



#%%