# Experiment 7: CycleGAN Horse2Zebra style transfer (TensorFlow)
# Requires: tensorflow, tensorflow-datasets, matplotlib | downloads Horse2Zebra (internet); GPU recommended
# Exact code of Exp7.ipynb. Notebook-only lines (!pip / %pip) are kept as comments so the file runs as normal Python.
# ===========================================================================

# !pip install tensorflow-datasets

# %pip install importlib-resources

import tensorflow as tf
import tensorflow_datasets as tsdf
import matplotlib.pyplot as plt

data = tsdf.load(
    "cycle_gan/horse2zebra",
    as_supervised=True
)

IMG_SIZE=256
EPOCHS=50
LAMBDA=10
BATCH_SIZE=1

print(data.keys())

def preprocess(image, label):
    image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
    image = tf.cast(image, tf.float32)
    image = image / 127.5 - 1.0

    return image

A = data["trainA"].map(preprocess).shuffle(100).batch(BATCH_SIZE)
B = data["trainB"].map(preprocess).shuffle(100).batch(BATCH_SIZE)

dataset = tf.data.Dataset.zip((A,B)).prefetch(tf.data.AUTOTUNE)

from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Conv2D, Conv2DTranspose, LeakyReLU, Input

G_AB = Sequential([
    Input(shape=(256, 256, 3)),
    Conv2D(32, 7, padding="same", activation="relu"),
    Conv2D(64, 3, strides=2, padding="same", activation="relu"),
    Conv2D(128, 3, strides=2, padding="same", activation="relu"),
    Conv2DTranspose(64, 3, strides=2, padding="same", activation="relu"),
    Conv2DTranspose(32, 3, strides=2, padding="same", activation="relu"),
    Conv2D(3, 7, padding="same", activation="tanh")
])

G_BA = Sequential([
    Input(shape=(256, 256, 3)),
    Conv2D(32, 7, padding="same", activation="relu"),
    Conv2D(64, 3, strides=2, padding="same", activation="relu"),
    Conv2D(128, 3, strides=2, padding="same", activation="relu"),
    Conv2DTranspose(64, 3, strides=2, padding="same", activation="relu"),
    Conv2DTranspose(32, 3, strides=2, padding="same", activation="relu"),
    Conv2D(3, 7, padding="same", activation="tanh")
])

D_A = Sequential([
    Input(shape=(256, 256, 3)),
    Conv2D(32, 4, strides=2, padding="same", activation="relu"),
    LeakyReLU(negative_slope=0.2),
    Conv2D(64, 4, strides=2, padding="same", activation="relu"),
    LeakyReLU(negative_slope=0.2),
    Conv2D(128, 4, strides=2, padding="same", activation="relu"),
    LeakyReLU(negative_slope=0.2),

    Conv2D(1, 4, padding="same")
])

D_B = Sequential([
    Input(shape=(256, 256, 3)),
    Conv2D(32, 4, strides=2, padding="same", activation="relu"),
    LeakyReLU(negative_slope=0.2),
    Conv2D(64, 4, strides=2, padding="same", activation="relu"),
    LeakyReLU(negative_slope=0.2),
    Conv2D(128, 4, strides=2, padding="same", activation="relu"),
    LeakyReLU(negative_slope=0.2),

    Conv2D(1, 4, padding="same")
])

loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)

g_ab_opt = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
g_ba_opt = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
d_a_opt = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
d_b_opt = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)

def cycle_loss(real, pred):
    return tf.reduce_mean(tf.abs(real-pred))

def gan_loss(target, pred):
    return loss_fn(target, pred)


@tf.function
def train_step(real_A, real_B):
    with tf.GradientTape(persistent=True) as tape:

        fake_B = G_AB(real_A, training=True)
        fake_A = G_BA(real_B, training=True)

        reconA = G_BA(fake_B, training=True)
        reconB = G_AB(fake_A, training=True)

        pred_B = D_B(fake_B, training=True)
        pred_A = D_A(fake_A, training=True)

        g_ab_loss = gan_loss(tf.ones_like(pred_B), pred_B)
        g_ba_loss = gan_loss(tf.ones_like(pred_A), pred_A)

        cyc_A = cycle_loss(real_A, reconA)
        cyc_B = cycle_loss(real_B, reconB)
        cyc = cyc_A + cyc_B

        tot_g_ab = g_ab_loss + LAMBDA*cyc
        tot_g_ba = g_ba_loss + LAMBDA*cyc

        pred_A_real = D_A(real_A, training=True)
        pred_B_real = D_B(real_B, training=True)
        fake_A_pred = D_A(tf.stop_gradient(fake_A), training=True)
        fake_B_pred = D_B(tf.stop_gradient(fake_B), training=True)

        d_a_loss = (
            gan_loss(tf.ones_like(pred_A_real), pred_A_real) + gan_loss(tf.zeros_like(fake_A_pred), fake_A_pred)
        ) / 2
        
        d_b_loss = (
            gan_loss(tf.ones_like(pred_B_real), pred_B_real) + gan_loss(tf.zeros_like(fake_B_pred), fake_B_pred)
        ) / 2

    g_ab_opt.apply_gradients(zip(
        tape.gradient(tot_g_ab, G_AB.trainable_variables),
        G_AB.trainable_variables
    ))
    g_ba_opt.apply_gradients(zip(
        tape.gradient(tot_g_ba, G_BA.trainable_variables),
        G_BA.trainable_variables
    ))

    d_a_opt.apply_gradients(zip(
        tape.gradient(d_a_loss, D_A.trainable_variables),
        D_A.trainable_variables
    ))
    d_b_opt.apply_gradients(zip(
        tape.gradient(d_b_loss, D_B.trainable_variables),
        D_B.trainable_variables
    ))

    return tot_g_ab+tot_g_ba, d_a_loss+d_b_loss


for epoch in range(EPOCHS):
    for real_A, real_B in dataset:
        g_loss, d_loss = train_step(real_A, real_B)

    print(
                f"Epoch {epoch+1}/{EPOCHS} | "
                f"G Loss: {g_loss.numpy():.4f} | "
                f"D Loss: {d_loss.numpy():.4f}"
            )
