# Experiment 1: Autoencoder + Variational Autoencoder on MNIST (TensorFlow/Keras)
# Requires: tensorflow, matplotlib, pandas | MNIST is downloaded by Keras on the first run (internet)
# Exact code of Exp1.ipynb. Notebook-only lines (!pip / %pip) are kept as comments so the file runs as normal Python.
# ===========================================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from tensorflow.keras.datasets import mnist
from tensorflow.keras import Sequential, Model
from tensorflow.keras.layers import Dense, Input, Lambda

import tensorflow as tf

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

X_test = X_test.astype("float32") / 255.0
X_train = X_train.astype("float32") / 255.0

X_test = X_test.reshape(-1, 784)
X_train = X_train.reshape(-1, 784)

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

for i in range(10):
    plt.subplot(2, 5, i+1)
    plt.imshow(X_train[i].reshape(28,28), cmap='gray')
    plt.title(f"Digit {y_train[i]}")
    plt.axis('off')

encoder = Sequential([
    Input(shape=(784, )),
    Dense(256, activation="relu"),
    Dense(128, activation="relu"),
    Dense(64, activation="relu"),
])

decoder = Sequential([  
    Input(shape=(64, )),
    Dense(128, activation="relu"),
    Dense(256, activation="relu"),
    Dense(784, activation="sigmoid"),
])

autoencoder = Sequential([
    encoder,
    decoder
])

autoencoder.summary()

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

history = autoencoder.fit(
    X_train, X_train,
    epochs=10,
    validation_data=(X_test, X_test),
    batch_size=32
)

plt.plot(history.history["loss"])
plt.plot(history.history["val_loss"])

reconstructed = autoencoder.predict(X_test[:10])

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

for i in range(10):
  plt.subplot(2, 10, i+1)
  plt.imshow(X_test[i].reshape(28, 28), cmap="gray")
  plt.axis("off")

  plt.subplot(2, 10, i+11)
  plt.imshow(reconstructed[i].reshape(28, 28), cmap="gray")
  plt.axis("off")

for i in range(10):
    t = encoder.predict(X_train[i].reshape(1, 784))
    print(f"Latent vector for {y_train[i]}: \n{t}")

(X_train, _), (X_test, _) = mnist.load_data()

X_train = X_train.astype("float32") / 255.0
X_test = X_test.astype("float32") / 255.0

X_train = X_train.reshape(-1, 784)
X_test = X_test.reshape(-1, 784)

inputs = Input(shape=(784, ))
h1 = Dense(512, activation="relu")(inputs)
h2 = Dense(256, activation="relu")(h1)
h3 = Dense(128, activation="relu")(h2)

z_mean = Dense(6)(h3)
z_log_var = Dense(6)(h3)

def sampling(args):
    mean, log_var = args
    epsilon = tf.random.normal(shape=tf.shape(mean))
    return mean + tf.exp(log_var*0.5)*epsilon

z = Lambda(sampling)([z_mean, z_log_var])

encoder = Model(inputs, [z_mean, z_log_var, z])
encoder.summary()

latent_input = Input(shape=(6,))
h1 = Dense(128, activation="relu")(latent_input)
h2 = Dense(256, activation="relu")(h1)
h3 = Dense(512, activation="relu")(h2)
outputs = Dense(784, activation="sigmoid")(h3)

decoder = Model(latent_input, outputs)

decoder.summary()

class VAE(Model):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder 

    def train_step(self, data):

        with tf.GradientTape() as tape:
            mean, log_var, z = self.encoder(data, training=True)
            recon = self.decoder(z, training=True)

            recon_loss = tf.reduce_mean(
                tf.reduce_sum(
                    tf.keras.backend.binary_crossentropy(data, recon),
                    axis=1
                )
            )

            kl_loss = -0.5 * tf.reduce_mean(
                tf.reduce_sum(
                    1 + log_var - tf.square(mean) - tf.exp(log_var),
                    axis=1
                )
            )

            loss = recon_loss + kl_loss

        grads = tape.gradient(loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))

        return {
            "loss": loss,
            "reconstruction_loss": recon_loss,
            "kl_loss": kl_loss
        }

vae = VAE(encoder, decoder)
vae.compile(tf.keras.optimizers.Adam(learning_rate=0.0001))

vae.fit(X_train, epochs=5, batch_size=128)

z_samples = np.random.normal(size=(10, 6)).astype("float32")
generated = decoder.predict(z_samples)

# 9. Display generated images
plt.figure(figsize=(12, 4))

for i in range(10):
    plt.subplot(2, 5, i + 1)
    plt.imshow(generated[i].reshape(28, 28), cmap="gray")
    plt.axis("off")

plt.show()
