import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.datasets import mnist
import matplotlib.pyplot as plt

# Load and preprocess the MNIST data
(x_train, _), (x_test, _) = mnist.load_data()  # Load the dataset and ignore the labels
x_train = x_train.astype("float32") / 255.0  # Normalize the training images to the range [0, 1]
x_test = x_test.astype("float32") / 255.0    # Normalize the test images to the range [0, 1]

# Reshape the data to be (samples, 784)
x_train = x_train.reshape((-1, 28 * 28))  # Flatten images to vectors of size 784
x_test = x_test.reshape((-1, 28 * 28))    # Flatten test images to vectors of size 784

# Add noise to the data
def add_noise(data, noise_factor=0.5):
    """
    Add Gaussian noise to the data.

    Parameters:
    data (numpy.ndarray): Input data (images).
    noise_factor (float): The amount of noise to add.

    Returns:
    numpy.ndarray: Noisy data with pixel values clipped to [0, 1].
    """
    noise = np.random.normal(loc=0.0, scale=1.0, size=data.shape)  # Generate Gaussian noise
    noisy_data = data + noise_factor * noise  # Add noise to the original data
    return np.clip(noisy_data, 0., 1.)  # Ensure pixel values remain in [0, 1]

# Generate noisy training and test data
x_train_noisy = add_noise(x_train)  # Add noise to the training data
x_test_noisy = add_noise(x_test)    # Add noise to the test data

# Define the autoencoder model
input_img = layers.Input(shape=(x_train.shape[1],))  # Input layer with shape (784,)
# Encoder: progressively reduce dimensionality
encoded = layers.Dense(200, activation="relu")(input_img)  # First hidden layer with 200 neurons
encoded = layers.Dense(100, activation="relu")(encoded)    # Second hidden layer with 100 neurons
encoded = layers.Dense(50, activation="relu")(encoded)     # Third hidden layer with 50 neurons
bottleneck = layers.Dense(25, activation="relu")(encoded)  # Bottleneck layer with 25 neurons

# Decoder: progressively increase dimensionality to reconstruct the input
decoded = layers.Dense(50, activation="relu")(bottleneck)   # First hidden layer in decoder
decoded = layers.Dense(100, activation="relu")(decoded)     # Second hidden layer in decoder
decoded = layers.Dense(200, activation="relu")(decoded)     # Third hidden layer in decoder
decoded_output = layers.Dense(x_train.shape[1], activation="sigmoid")(decoded)  # Output layer with 784 neurons

# Create the autoencoder model
autoencoder = models.Model(input_img, decoded_output)  # Define the model with input and output
autoencoder.compile(optimizer="adam", loss="mse")      # Compile the model with Adam optimizer and MSE loss

# Print the model summary
print("Autoencoder Model Summary:")
autoencoder.summary()  # Display the model architecture

# Train the autoencoder
autoencoder.fit(x_train_noisy, x_train, epochs=40, batch_size=256, validation_data=(x_test_noisy, x_test))  
# Train the model with noisy images as inputs and original images as targets

# Evaluate the model by denoising the test images
denoised_images = autoencoder.predict(x_test_noisy)  # Use the trained model to denoise the test images

# Visualization of results
n = 10  # Number of images to display
plt.figure(figsize=(20, 6))

# Original Images
for i in range(n):
    # Display original images
    ax = plt.subplot(3, n, i + 1)  # Create a subplot for the original image
    plt.imshow(x_test[i].reshape(28, 28), cmap="gray")  # Display original image
    plt.title("Original Image")  # Set title for the subplot
    plt.axis("off")  # Hide axis

    # Display noisy images
    ax = plt.subplot(3, n, i + 1 + n)  # Create a subplot for the noisy image
    plt.imshow(x_test_noisy[i].reshape(28, 28), cmap="gray")  # Display noisy image
    plt.title("Noisy Image")  # Set title for the subplot
    plt.axis("off")  # Hide axis

    # Display denoised images
    ax = plt.subplot(3, n, i + 1 + 2 * n)  # Create a subplot for the denoised image
    plt.imshow(denoised_images[i].reshape(28, 28), cmap="gray")  # Display denoised image
    plt.title("Denoised Image")  # Set title for the subplot
    plt.axis("off")  # Hide axis

plt.tight_layout()  # Adjust layout to prevent overlap
plt.show()  # Display the plots