# Experiment 5: Autoencoder DeepFake: shared encoder + two decoders (MNIST 6 <-> 8)
# Requires: tensorflow, matplotlib, pandas, scikit-learn | MNIST download (internet)
# Exact code of Exp5.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

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

from sklearn.model_selection import train_test_split

(X_train, y_train), (X_test, y_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))

numA = []
numB = []

for i in range(len(X_train)):
    if(y_train[i] == 6):
        numA.append(X_train[i])
    elif(y_train[i] == 8):
        numB.append(X_train[i])
    else: 
        continue

numA = np.array(numA, dtype="float32")
numB = np.array(numB, dtype="float32")

numA, numA_test = train_test_split(numA, train_size=0.8, random_state=42)
numB, numB_test = train_test_split(numB, train_size=0.8, random_state=42)

plt.figure(figsize=(12,4))
for i in range(10):
    plt.subplot(2, 5,i+1)
    plt.imshow(numA[i].reshape(28,28), cmap="grey")
    plt.axis("off")

plt.figure(figsize=(12,4))
for i in range(10):
    plt.subplot(2, 5,i+1)
    plt.imshow(numB[i].reshape(28,28), cmap="grey")
    plt.axis("off")

encoder = Sequential([
    Input(shape=(784, )),
    Dense(512, activation="relu"),
    Dense(256, activation="relu"),
    Dense(128, activation="relu"),
    Dense(20, activation="relu", name="latent_rep")
])

encoder.summary()

decoderA = Sequential([
    Input(shape=(20, )),
    Dense(128, activation="relu"),
    Dense(256, activation="relu"),
    Dense(512, activation="relu"),
    Dense(784, activation="sigmoid"),
])

decoderB = Sequential([
    Input(shape=(20, )),
    Dense(128, activation="relu"),
    Dense(256, activation="relu"),
    Dense(512, activation="relu"),
    Dense(784, activation="sigmoid"),
])

decoderA.summary()

autoencoderA = Sequential([encoder, decoderA])
autoencoderB = Sequential([encoder, decoderB])

autoencoderA.compile(optimizer="adam", loss="binary_crossentropy")
autoencoderB.compile(optimizer="adam", loss="binary_crossentropy")

autoencoderA.fit(numA, numA, epochs=10, batch_size=32)

autoencoderB.fit(numB, numB, epochs=10, batch_size=32)

#a->a

reconA = autoencoderA.predict(numA_test[:5])

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

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

    plt.subplot(2, 5, i+6)
    plt.imshow(reconA[i].reshape(28, 28), cmap="grey")
    plt.axis("off")

#b->b

reconB = autoencoderB.predict(numB_test[:5])

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

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

    plt.subplot(2, 5, i+6)
    plt.imshow(reconB[i].reshape(28, 28), cmap="grey")
    plt.axis("off")

#a->b

latentA = encoder.predict(numA_test[:10])
reconAB = decoderB.predict(latentA)

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

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

    plt.subplot(2, 5, i+6)
    plt.imshow(reconAB[i].reshape(28, 28), cmap="grey")
    plt.axis("off")

#b->a

latentB = encoder.predict(numB_test[:10])
reconBA = decoderA.predict(latentB)

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

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

    plt.subplot(2, 5, i+6)
    plt.imshow(reconBA[i].reshape(28, 28), cmap="grey")
    plt.axis("off")
