#Import necessary libraries:

# Import TensorFlow and NumPy
import tensorflow as tf
import numpy as np

# Import the train and test split function from sklearn:

from sklearn.model_selection import train_test_split

# Take only 10% of the data:

# Extract a 10% sample from the training dataset
x_train, _, y_train, _ = train_test_split(x_train_total, 
                                          y_train_total, 
                                          test_size=0.9, 
                                          random_state=1502)
 
# Extract a 10% sample from the test dataset
x_test, _, y_test, _ = train_test_split(x_test_total, 
                                        y_test_total, 
                                        test_size=0.9, 
                                        random_state=1502)

# Define the class labels of the Fashion MNIST dataset

# Define the class names corresponding to the labels in the dataset
class_names = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
               "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]

# Normalize the training and testing datasets:

# Normalize the pixel values to the range [0,1]
x_train = x_train / 255.0
x_test = x_test / 255.0

# Adjust the shape of datasets to encompass the channel dimension:

# Alter the data shape to be (number of samples, height, width, number of channels)
x_train = x_train.reshape((-1, 28, 28, 1))
x_test = x_test.reshape((-1, 28, 28, 1))

# Set up the ImageDataGenerator using the given parameters to alter the image data as specified.

# Initializing the ImageDataGenerator with the desired augmentations
datagen = tf.keras.preprocessing.image.ImageDataGenerator(
   rotation_range=10,      # Rotate images up to 10 degrees
   zoom_range=0.1,         # Zoom in on images up to 10%
   shear_range=0.1,        # Shear images by a factor of 0.1
   width_shift_range=0.1,  # Shift the width of images up to 10%
   height_shift_range=0.1, # Shift the height of images up to 10%
   horizontal_flip=True    # Allow horizontal flipping of images
)

# Implement the defined augmentations to the main training dataset.

# Apply the defined data augmentations to the training set
datagen.fit(x_train)

# Start by importing the necessary TensorFlow libraries.

# Importing required modules from TensorFlow
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout

# Construct the CNN using the specified architecture.

# Defining the CNN architecture with the provided specifications
model = Sequential([
    Conv2D(32, (3,3), activation='relu', input_shape=(28, 28, 1)),  # First convolutional layer
    MaxPooling2D(2, 2),                                             # First pooling layer
    Dropout(0.2),                                                   # 20% dropout layer
    Conv2D(64, (3,3), activation='relu'),                           # Second convolutional layer
    MaxPooling2D(2, 2),                                             # Second pooling layer
    Dropout(0.2),                                                   # 20% dropout layer
    Flatten(),                                                      # Flattening layer
    Dense(256, activation='relu'),                                  # Dense layer with 256 units
    Dropout(0.5),                                                   # 50% dropout layer
    Dense(10, activation='softmax')                                 # Output layer with 10 units
])

# Compile the CNN model with the appropriate parameters.

# Compiling the model with the Adam optimizer and suitable loss function
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Set up the Early Stopping callback.

# Implementing early stopping based on validation accuracy
early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_accuracy',
                                              patience=5)

# Define the custom learning rate scheduler function.

# Defining the learning rate schedule function
def lr_schedule(epoch, lr):
    if epoch > 10:
        return lr * 0.1  # Reduce learning rate by a factor of 0.1 after 10 epochs
    return lr

# Initialize the Learning Rate Scheduler callback using the custom function.

# Initializing the learning rate scheduler callback
scheduler = tf.keras.callbacks.LearningRateScheduler(lr_schedule)

# Start by importing necessary visualization and matrix evaluation libraries.

from sklearn.metrics import confusion_matrix
import seaborn as sns

# Start by naming the function, the inputs and writing the documentation.
# Acquire the desired convolutional layer from the model and construct a new model that maps the original input to this layer's activations.
# Determine the index or position of our target layer within the model's architecture.
# Establish a separate model that processes the activations of our target layer and yields the model's ultimate class predictions.
# Utilize TensorFlow's gradient tracking mechanism to monitor the activations and compute the classification scores.
# Compute the gradient that represents the importance of each feature in the convolutional layer's activations with respect to the target class score.
# Generate the heatmap by weighing the activations by their respective importance, and then normalize it to ensure values lie between 0 and 1.


def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    """
    Generate a GradCAM heatmap for a given image array.
    
    Args:
    - img_array (numpy.ndarray): The input image array.
    - model (tf.keras.Model): The pre-trained model.
    - last_conv_layer_name (str): The name of the last convolutional layer in the model.
    - pred_index (int, optional): The index of the prediction class. Uses top prediction if None.
    
    Returns:
    - numpy.ndarray: The GradCAM heatmap.
    """
    # Obtain the last convolutional layer from the model
    last_conv_layer = model.get_layer(last_conv_layer_name)
 
    # Create a model that maps the model's input to the last_conv_layer's output
    last_conv_layer_model = tf.keras.Model(model.inputs, last_conv_layer.output)
 
 
    # Get the index of the specified convolutional layer in the model
    last_conv_layer_idx = [layer.name for layer in model.layers].index(last_conv_layer_name)
 
 
    # Create an input placeholder with the shape of the output from last_conv_layer
    classifier_input = tf.keras.Input(shape=last_conv_layer.output.shape[1:])
    x = classifier_input
 
    # Loop through layers after the last_conv_layer and build the classifier model
    for layer in model.layers[last_conv_layer_idx+1:]:
        x = layer(x)
    classifier_model = tf.keras.Model(classifier_input, x)
 
    # Set up the GradientTape for automatic differentiation
    with tf.GradientTape() as tape:
        # Compute the activations of the last convolutional layer
        last_conv_layer_output = last_conv_layer_model(img_array)
        # Make sure TensorFlow watches this output for gradient calculations
        tape.watch(last_conv_layer_output)
        # Predict the class scores using the classifier model
        preds = classifier_model(last_conv_layer_output)
        # If no specific class index is provided, use the class with the highest score
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        # Extract the score for the target class
        class_channel = preds[:, pred_index]
 
    # Calculate the gradient of the target class score with respect to the activations
    grads = tape.gradient(class_channel, last_conv_layer_output)
 
    # Average the gradients spatially (over height and width) to get the importance of each feature map
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
 
 
    # Weigh the activations by the importance to get the heatmap
    last_conv_layer_output = last_conv_layer_output * pooled_grads[tf.newaxis, tf.newaxis, tf.newaxis, :]
 
    # Average over the feature maps to get a 2D heatmap
    heatmap = tf.reduce_mean(last_conv_layer_output, axis=-1)[0]
    
    # Normalize the heatmap to lie between 0 and 1
    heatmap = np.maximum(heatmap, 0) / np.max(heatmap)
 
    return heatmap

# Randomly pick an image from x_test.

# Select a random image from the test dataset
img_array = x_test[np.random.choice(x_test.shape[0])]
 
# Reshape the image to add an extra dimension for batch size
img_array = img_array[np.newaxis, :]

# Use the make_gradcam_heatmap function to produce the heatmap.

# Generate the GradCAM heatmap for the image based on the specified layer and prediction
heatmap = make_gradcam_heatmap(img_array, model, 'conv2d', top_pred_index)

# Begin by saving the model in the native Keras format.

# Save the model
model.save("cnn_model.keras")

# Load the saved model back into the environment.

# Load the model
loaded_model = tf.keras.models.load_model("cnn_model.keras")

# Randomly select an image from the test dataset x_test.

# Predict on a random test image
random_idx = np.random.choice(len(x_test))
test_img = x_test[random_idx]