import numpy as np
from tensorflow.keras.datasets import mnist

# Load the MNIST dataset
(X_train, y_train), (X_test, y_test) = mnist.load_data()

# 1. Normalize the images by dividing by the maximum pixel value (255), to scale pixel values to the range [0, 1].
X_train = X_train.astype('float32') / 255.0
X_test = X_test.astype('float32') / 255.0

# 2. Reshape the images from 28x28 to 1D arrays of size 784 (flattening)
X_train = X_train.reshape((X_train.shape[0], 784))
X_test = X_test.reshape((X_test.shape[0], 784))

# Optional: Add noise to the dataset for denoising autoencoders
noise_factor = 0.5  # Adjust the noise factor for different levels of noise
X_train_noisy = X_train + noise_factor * np.random.normal(loc=0.0, scale=1.0, size=X_train.shape)
X_test_noisy = X_test + noise_factor * np.random.normal(loc=0.0, scale=1.0, size=X_test.shape)

# Clip the values to ensure they remain in the range [0, 1]
X_train_noisy = np.clip(X_train_noisy, 0.0, 1.0)
X_test_noisy = np.clip(X_test_noisy, 0.0, 1.0)

# The data is now preprocessed and ready for model training.
print("Preprocessing completed.")






# Q1: Load MNIST dataset and split it into training and test sets

from tensorflow.keras.datasets import mnist

import numpy as np

import matplotlib.pyplot as plt

import tensorflow as tf

(X_train, y_train), (X_test, y_test) = mnist.load_data()

# Q2: Normalize the data by dividing by the max pixel value (255)

X_train = X_train.astype('float32') / 255.

X_test = X_test.astype('float32') / 255.

# Q3: Print the shape of the dataset

print("X_train shape:", X_train.shape)

print("X_test shape:", X_test.shape)

# Q4: Reshape the images into 2D arrays (original, 784)

X_train = X_train.reshape(-1, 784)

X_test = X_test.reshape(-1, 784)

print("Reshaped X_train shape:", X_train.shape)

print("Reshaped X_test shape:", X_test.shape)

# Q5: Define the size of the encodings to 32

encoding_dim = 32

# Q6: Define the placeholders for the input (in Keras, input layers are defined via Input)

input_img = tf.keras.layers.Input(shape=(784,))

# Q7: Define the encoder which takes input_image and returns the encodings

encoded = tf.keras.layers.Dense(encoding_dim, activation='relu')(input_img)

# Q8: Define the decoder which reconstructs the image from the encoding

decoded = tf.keras.layers.Dense(784, activation='sigmoid')(encoded)

# Q9: Build the model (encoder + decoder)

autoencoder = tf.keras.Model(input_img, decoded)

# Q10: Find the number of parameters in the model

autoencoder.summary()

# Q11: Compile the model with an appropriate optimizer and loss function

autoencoder.compile(optimizer='adam', loss='binary_crossentropy')

# Q12: Fit the model for 10 epochs

autoencoder.fit(X_train, X_train,

                epochs=10,

                batch_size=256,

                shuffle=True,

                validation_data=(X_test, X_test))

# Q13: Reconstruct and plot the images

decoded_imgs = autoencoder.predict(X_test)

# Plot original vs reconstructed

n = 10

plt.figure(figsize=(20, 4))

for i in range(n):

    # display original

    ax = plt.subplot(2, n, i + 1)

    plt.imshow(X_test[i].reshape(28, 28), cmap='gray')

    ax.get_xaxis().set_visible(False)

    ax.get_yaxis().set_visible(False)

    

    # display reconstruction

    ax = plt.subplot(2, n, i + 1 + n)

    plt.imshow(decoded_imgs[i].reshape(28, 28), cmap='gray')

    ax.get_xaxis().set_visible(False)

    ax.get_yaxis().set_visible(False)

plt.show()

# Q14: Add noise to the train and test images and plot

noise_factor = 1.0

X_train_noisy = X_train + noise_factor * np.random.normal(loc=0.0, scale=1.0, size=X_train.shape)

X_test_noisy = X_test + noise_factor * np.random.normal(loc=0.0, scale=1.0, size=X_test.shape)

X_train_noisy = np.clip(X_train_noisy, 0., 1.)

X_test_noisy = np.clip(X_test_noisy, 0., 1.)

# Train on noisy data

autoencoder.fit(X_train_noisy, X_train,

                epochs=10,

                batch_size=256,

                shuffle=True,

                validation_data=(X_test_noisy, X_test))

# Reconstruct noisy test images

decoded_noisy_imgs = autoencoder.predict(X_test_noisy)

# Plot original vs noisy reconstructed

plt.figure(figsize=(20, 4))

for i in range(n):

    # display noisy images

    ax = plt.subplot(2, n, i + 1)

    plt.imshow(X_test_noisy[i].reshape(28, 28), cmap='gray')

    ax.get_xaxis().set_visible(False)

    ax.get_yaxis().set_visible(False)

    

    # display noisy reconstructions

    ax = plt.subplot(2, n, i + 1 + n)

    plt.imshow(decoded_noisy_imgs[i].reshape(28, 28), cmap='gray')

    ax.get_xaxis().set_visible(False)

    ax.get_yaxis().set_visible(False)

plt.show()