# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="S8alK1zLkbXK" colab_type="text"
# ## Introduction
#
# The objective is to use Deep Convolutional Generative Adversertial Network to generate fake GO game snapshots on a GO board.

# + [markdown] id="iMTrlOg3nz6u" colab_type="text"
# ## Import Libs

# + id="8588hUR4rzY7" colab_type="code" outputId="870aad12-16de-47a6-8b06-947c9f51d226" colab={"base_uri": "https://localhost:8080/", "height": 34}
from google.colab import drive
from keras.datasets import mnist
from keras.layers import Input, Dense, Reshape, Flatten, Dropout
from keras.layers import BatchNormalization, Activation, ZeroPadding2D
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import UpSampling2D, Conv2D
from keras.models import Sequential, Model
from keras.optimizers import Adam
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
import numpy as np
import time
import os
from PIL import Image
from PIL import Image
import numpy as np
import pandas as pd

# + [markdown] id="MKBG0PiCn9po" colab_type="text"
# ## Download Data
#
# * Data prepared offline by running extra python scripts from the listed repository (which is also created for this project)
#
# * Please reference slides

# + id="qQ7ECrvx1cnE" colab_type="code" outputId="9a892dfc-de4d-4253-d8c3-22c47d1a8a70" colab={"base_uri": "https://localhost:8080/", "height": 119}
# !rm -rf gangengo
# !git clone https://github.com/timlyrics/gangengo
# !mv gangengo/* .

# + [markdown] id="PyA6tSkalZ2Y" colab_type="text"
# ## Real sample

# + id="klChc48A2Ln3" colab_type="code" outputId="abb3e6e6-fc7e-471d-f0d4-9fbfc6c1c942" colab={"base_uri": "https://localhost:8080/", "height": 350}
from gen_go import GenGo
gg = GenGo()
gg.show_random()

# + [markdown] _uuid="04304e5c4555849eade69a1a1b5b1fb7cc2feb13" id="AGzaeNUNxdGy" colab_type="text"
# ## Set up network parameters
#
# * The go board is 19x19, but we padded it to 20x20 for easier processing.
#
# * Each pixel has two channels for black and white (red).
#
# * Generator latent dimension is 32.
#  

# + _uuid="38443112fd55bef54d7807495b69d66353bfccae" id="YCSOpXVVxdG1" colab_type="code" colab={}
img_rows = 20
img_cols = img_rows
channels = 2
img_shape = (img_rows, img_cols, channels)
latent_dim = 32


# + [markdown] _uuid="1c5ad4f2b1b3b4ae3dd6caf7ba96ef1e273eb900" id="McWm3ToexdG4" colab_type="text"
# ## Define a function to build a generator
#
# * Generator is a sequential model with one dense layer followed by two upsampling + convolution layers to buff complexity and match dimension.
#
# * Then a convolutional layer with 2 channels is used to match pixel depth and a sigmoid activation converts output to  boolean.

# + _uuid="804d0ab8433b786150c48cd6bd990d74705cb2a8" id="yIh80aD4xdG5" colab_type="code" colab={}
def build_generator():
    model = Sequential()
    model.add(Dense(128 * 5 * 5, activation="relu", input_dim=latent_dim))
    model.add(Reshape((5, 5, 128)))
    model.add(UpSampling2D())
    model.add(Conv2D(128, kernel_size=3, padding="same"))
    model.add(BatchNormalization(momentum=0.8))
    model.add(Activation("relu"))
    model.add(UpSampling2D())
    model.add(Conv2D(64, kernel_size=3, padding="same"))
    model.add(BatchNormalization(momentum=0.8))
    model.add(Activation("relu"))
    model.add(Conv2D(channels, kernel_size=3, padding="same"))
    model.add(Activation("sigmoid"))
    model.summary()
    noise = Input(shape=(latent_dim,))
    img = model(noise)
    m = Model(noise, img)
    return m


# + [markdown] _uuid="a9ea4efecf7d4faeb0a37a649c079ebe88396908" id="IP63a3hixdG7" colab_type="text"
# ## Define a function to build a discriminator
#
# * Discriminator a typical sequential Convolutional Neural Network classifier.
#
# * First layer padding is "valid" because the GO game treats sides and corders differently.

# + _uuid="b9aa34a385dac406b50e7e122b45483c98c0cffa" id="YSIwur4txdG8" colab_type="code" colab={}
def build_discriminator():
    model = Sequential()
    model.add(Conv2D(32, kernel_size=3, strides=2, input_shape=img_shape, padding="valid"))
    model.add(LeakyReLU(alpha=0.2))
    model.add(Dropout(0.25))
    model.add(Conv2D(64, kernel_size=3, strides=2, padding="same"))
    model.add(ZeroPadding2D(padding=((0, 1), (0, 1))))
    model.add(BatchNormalization(momentum=0.8))
    model.add(LeakyReLU(alpha=0.2))
    model.add(Dropout(0.25))
    model.add(Conv2D(128, kernel_size=3, strides=2, padding="same"))
    model.add(BatchNormalization(momentum=0.8))
    model.add(LeakyReLU(alpha=0.2))
    model.add(Dropout(0.25))
    model.add(Conv2D(256, kernel_size=3, strides=1, padding="same"))
    model.add(BatchNormalization(momentum=0.8))
    model.add(LeakyReLU(alpha=0.2))
    model.add(Dropout(0.25))
    model.add(Flatten())
    model.add(Dense(1, activation='sigmoid'))
    model.summary()
    img = Input(shape=img_shape)
    validity = model(img)
    return Model(img, validity)


# + [markdown] _uuid="306fda1d4980efc3c64db997111ab7e615cfdbac" id="v3U8x5AZxdG-" colab_type="text"
# ## Build GAN
#
# * The discriminator model is a simple typical classification model.
#
# * The combined model is a the discriminator stacked on top of generator where discriminator parameters are freezed.
#
# * Binary cross entropy is used for loss for both discriminator loss and combined loss.
#
# * Adam optimizer is used for both models where learning rate for combined model is 0.0001 and for discriminator is 0.00002 (1 order of magnitude lower than generator) and a decay is also added to the discriminator optimizer to prevent overpowering the generator.
#
# (The parameters are tuned by trial and error until time ran out)

# + _uuid="f7202cc45aae945674176146922493812e520030" id="c1RvpIcoxdG_" colab_type="code" outputId="17250a9b-8c5f-4649-a78d-6a6fd17cf761" colab={"base_uri": "https://localhost:8080/", "height": 1397}
# build discriminator
discriminator = build_discriminator()
discriminator.compile(loss='binary_crossentropy',
                      optimizer=Adam(0.00002, 0.3, decay=1e-6),
                      metrics=['accuracy'])

# build generator
generator = build_generator()
z = Input(shape=(latent_dim,))
img = generator(z)

# For the combined model we will only train the generator
discriminator.trainable = False

# The discriminator takes generated images as input and determines validity
valid = discriminator(img)

# The combined model  (stacked generator and discriminator)
# Trains the generator to fool the discriminator
combined = Model(z, valid)
combined.compile(loss='binary_crossentropy', optimizer=Adam(0.0001, 0.9))


# + [markdown] id="0JqoNWtXh6PP" colab_type="text"
# ## Transform between softmax space and linear space

# + id="q7bXdNBqi3mb" colab_type="code" colab={}
def softmax_from_linear(arr):
  arr = np.expand_dims(arr, axis=3)
  smx = np.concatenate([1 * (arr > 0), 1 * (arr < 0)], axis=3)
  return smx

def linear_from_softmax(smx):
  arr = ((smx[:,:,:,0] > 0.5) & (smx[:,:,:,0] > smx[:,:,:,1])) * 1 + ((smx[:,:,:,1] > 0.5) & (smx[:,:,:,1] > smx[:,:,:,0])) * -1
  return arr


# + [markdown] _uuid="e1d9927bd8601eedf144b6f46f12cc4da0d045b5" id="8a5yO3JfxdHB" colab_type="text"
# ## Define a function to train GAN
#
# * Trains both models in each batch / epoch.
#
# * Saves images and records results periodically.

# + _uuid="6592840f1e019f1d6fd3e709a0885389d252164f" id="Fkhx_GCHxdHC" colab_type="code" colab={}
def train(epochs, batch_size, save_interval):
    os.makedirs('images', exist_ok=True)

    # Adversarial ground truths
    valid = np.ones((batch_size, 1))
    fake = np.zeros((batch_size, 1))

    # losses
    arr_d_loss = np.zeros(epochs)
    arr_g_loss = np.zeros(epochs)
    
    for epoch in range(epochs):
        # Select a random real images
        real_imgs = softmax_from_linear(gg.arr_from_random(batch_size))

        # Sample noise and generate a batch of fake images
        noise = np.random.normal(0, 1, (batch_size, latent_dim))
        fake_imgs = generator.predict(noise)
        
        # Train the discriminator
        D_loss_real = discriminator.train_on_batch(real_imgs, valid)
        D_loss_fake = discriminator.train_on_batch(fake_imgs, fake)
        D_loss = 0.5 * np.add(D_loss_real, D_loss_fake)

        # Train the generator
        g_loss = combined.train_on_batch(noise, valid)
        
        # Remember loss
        arr_d_loss[epoch] = D_loss[0]
        arr_g_loss[epoch] = float(g_loss)

        # If at save interval
        if epoch % save_interval == 0:
            # Print the progress
            print("%d [D loss: %f, acc.: %.2f%%] [G loss: %f]" % (epoch, D_loss[0], 100 * D_loss[1], g_loss))
            # Save generated image samples
            save_imgs(epoch)

    return arr_d_loss, arr_g_loss


# + [markdown] id="B6YIRgYfxXsW" colab_type="text"
# ## Define visualizer
#
# * Plot randomly generated samples in 3x3 layout

# + _uuid="b03c0413ddcde5125fa8fdb22bfe47e1bd668606" id="145zyE3WxdHE" colab_type="code" colab={}
def save_imgs(epoch):
    r, c = 3, 3
    noise = np.random.normal(0, 1, (r * c, latent_dim))
    gen_imgs = linear_from_softmax(generator.predict(noise))

    fig, axs = plt.subplots(r, c)
    fig.set_figheight(10)
    fig.set_figwidth(10)

    cnt = 0
    for i in range(r):
        for j in range(c):
            row = gen_imgs[cnt, :, :]
            cnt += 1
            
            ax = axs[i, j]
            loc = MultipleLocator(base=1)
            ax.xaxis.set_major_locator(loc)
            ax.yaxis.set_major_locator(loc)
            ax.set_axisbelow(True)
            ax.grid(linestyle='-', axis='both', linewidth='0.5', color='grey')
            for x in range(row.shape[0]):
                for y in range(row.shape[1]):
                    if row[x][y] > 0.5:
                        ax.scatter(x, y, c='red', s=50)
                    elif row[x][y] < -0.5:
                        ax.scatter(x, y, c='black', s=50)

    fig.savefig("images/go_%d.png" % epoch)
    plt.close()


# + [markdown] _uuid="ccbe0a280b9fc16e47a63dc5e5edae8909f7569c" id="2CQ4LNlaxdHJ" colab_type="text"
# ## Train GAN
#
# * Run training routine

# + _uuid="8a52d992f11a034e0cc0cabaddeb2ae93f160dc6" id="RvwGlSuMxdHK" colab_type="code" outputId="e6fefc8f-6095-47e9-9ce2-6409cb9dc01c" colab={"base_uri": "https://localhost:8080/", "height": 3556}
start = time.time()

arr_d_loss, arr_g_loss = train(epochs=100001, batch_size=128, save_interval=500)

end = time.time()
elapsed_train_time = 'elapsed training time: {} min, {} sec '.format(int((end - start) / 60),
                                                                     int((end - start) % 60))
print(elapsed_train_time)

# + [markdown] id="PLUGoE8sh0aE" colab_type="text"
# ## Loss evolution over time
#
# * The progression of the training proceeded well (combined loss decreases) until 20000 batches / epochs.
#
# * The discriminator overpowers the generator after 20000, and the training does not progress well.
#
# * Further tuning is necessary to undercut the discriminator to achieve better results.

# + id="Wn4smoAZOQqe" colab_type="code" outputId="e2d712fa-3cf5-4feb-eae3-446ee16ca392" colab={"base_uri": "https://localhost:8080/", "height": 378}
plt.plot(np.arange(1001) * 100, arr_d_loss[::100])
plt.plot(np.arange(1001) * 100, arr_g_loss[::100])
plt.legend(['discriminator loss', 'generator loss'])
plt.xlabel('epoch')
plt.ylabel('loss')

# + [markdown] id="zwFvb8tFirwI" colab_type="text"
# ## Save weights and do some debugging checks

# + _uuid="cb7c9e65347eb077e1dabd74445713ddef12ba34" id="5ih-0phwxdHO" colab_type="code" colab={}
os.makedirs('saved_model_weights', exist_ok=True)
generator.save_weights('saved_model_weights/generator_weights.h5')
discriminator.save_weights('saved_model_weights/discriminator_weights.h5')
combined.save_weights('saved_model_weights/combined_weights.h5')

# + id="WsHG1dqL0zqz" colab_type="code" colab={}
# debugging
noise = np.random.normal(0, 1, (1, latent_dim))
raw = generator.predict(noise)
gen_imgs = linear_from_softmax(generator.predict(noise))

# + colab_type="code" id="hnVgPu_rC9Gx" colab={}
# debugging
a = raw[0,:,:]
(a[:,:,0] < 0.5) & (a[:,:,1] < 0.5)

b = softmax_from_linear(gg.arr_from_random(1))

# + [markdown] _uuid="526455382f869ceadc4206cb703e676a4037b213" id="VY40mkpTxdHQ" colab_type="text"
# ## Show generated GO
#
# * Display randomly generated images for epoch 0, 500, 5000, 20000, and 100000.

# + _uuid="b9cd883f622680511e19cfa30717a3a8d93c4a48" id="dXDYEzoNxdHR" colab_type="code" outputId="444b93b6-4ccd-463b-aa7b-8538e2e26c5d" colab={"base_uri": "https://localhost:8080/", "height": 737}
Image.open('images/go_0.png')

# + _uuid="4c487c7fede8db3baa681554c9cf677464af791c" id="ON5zWZMoxdHW" colab_type="code" outputId="d10fe14d-c654-4410-a5ad-98f913fc1b3d" colab={"base_uri": "https://localhost:8080/", "height": 737}
Image.open('images/go_500.png')

# + id="fu4LQ12OlyjA" colab_type="code" outputId="4538e6ed-bfb3-4a4a-b207-4ed370839f63" colab={"base_uri": "https://localhost:8080/", "height": 737}
Image.open('images/go_5000.png')

# + id="OPeqA0s3l4SF" colab_type="code" outputId="742c10c9-b66b-4b11-f798-10d50dc408f8" colab={"base_uri": "https://localhost:8080/", "height": 737}
Image.open('images/go_20000.png')

# + id="EmTBu61qmAR9" colab_type="code" outputId="48545dce-6dd7-4d69-e7c9-69e527c89896" colab={"base_uri": "https://localhost:8080/", "height": 737}
Image.open('images/go_50000.png')

# + id="DPN14p71mGix" colab_type="code" outputId="8d9a91f5-92ab-4bb1-d59d-152315ec9062" colab={"base_uri": "https://localhost:8080/", "height": 737}
Image.open('images/go_100000.png')

# + [markdown] _uuid="e1a1fdfff62338dd332aa823655f398156a5e5ae" id="pqR-k48QxdHZ" colab_type="text"
# ## Reference
# [Keras - DCGAN](https://github.com/eriklindernoren/Keras-GAN#dcgan)

StarcoderdataJupyterScriptsDedupFiltered