# Path: utils/utils.py
# def put_kernels_on_grid(kernel, (grid_Y, grid_X), pad=1):
#     '''Visualize conv. features as an image (mostly for the 1st layer).
#     Place kernel into a grid, with some paddings between adjacent filters.
#     Args:
#       kernel:            tensor of shape [Y, X, NumChannels, NumKernels]
#       (grid_Y, grid_X):  shape of the grid. Require: NumKernels == grid_Y * grid_X
#                            User is responsible of how to break into two multiples.
#       pad:               number of black pixels around each filter (between them)
# 
#     Return:
#       Tensor of shape [(Y+pad)*grid_Y, (X+pad)*grid_X, NumChannels, 1].
#     '''
#     # pad X and Y
#     x1 = tf.pad(kernel, tf.constant([[pad, 0], [pad, 0], [0, 0], [0, 0]]))
# 
#     # X and Y dimensions, w.r.t. padding
#     Y = kernel.get_shape()[0] + pad
#     X = kernel.get_shape()[1] + pad
#     ch = kernel.get_shape()[2]
# 
#     # put NumKernels to the 1st dimension
#     x2 = tf.transpose(x1, (3, 0, 1, 2))
#     # organize grid on Y axis
#     x3 = tf.reshape(x2, tf.pack([grid_X, Y * grid_Y, X, ch]))
# 
#     # switch X and Y axes
#     x4 = tf.transpose(x3, (0, 2, 1, 3))
#     # organize grid on X axis
#     x5 = tf.reshape(x4, tf.pack([1, X * grid_X, Y * grid_Y, ch]))
# 
#     # back to normal order (not combining with the next step for clarity)
#     x6 = tf.transpose(x5, (2, 1, 3, 0))
# 
#     # to tf.image_summary order [batch_size, height, width, channels],
#     #   where in this case batch_size == 1
#     x7 = tf.transpose(x6, (3, 0, 1, 2))
# 
#     # scale to [0, 1]
#     x_min = tf.reduce_min(x7)
#     x_max = tf.reduce_max(x7)
#     x8 = (x7 - x_min) / (x_max - x_min)
# 
#     return x8


import os
import numpy as np
import tensorflow as tf
from keras import backend as K
from keras.layers import Input, Dense, Reshape, Activation, Convolution2D, Deconvolution2D, LeakyReLU, Flatten, \
    BatchNormalization as BN
from keras.models import Sequential, Model
from keras import initializations
from functools import partial
from utils import cifar
from utils.utils import put_kernels_on_grid


def get_model(sess, image_shape=(32, 32, 3), gf_dim=64, df_dim=64, batch_size=64,
              name="autoencoder"):
    K.set_session(sess)
    with tf.variable_scope(name):
        # sizes
        ch = image_shape[2]
        rows = [4, 8, 16, 32]
        cols = [4, 8, 16, 32]

        # nets
        G = generator(batch_size, gf_dim, ch, rows, cols)
        G.compile("sgd", "mse")
        g_vars = G.trainable_weights
        print "G.shape: ", G.output_shape

        E = encoder(batch_size, df_dim, ch, rows, cols)
        E.compile("sgd", "mse")
        e_vars = E.trainable_weights
        print "E.shape: ", E.output_shape

        D = discriminator(batch_size, df_dim, ch, rows, cols)
        D.compile("sgd", "mse")
        d_vars = D.trainable_weights
        print "D.shape: ", D.output_shape

        Z2 = Input(batch_shape=(batch_size, z_dim), name='more_noise')
        Z = G.input
        Img = D.input

# Path: utils/utils.py
# def put_kernels_on_grid(kernel, (grid_Y, grid_X), pad=1):
#     '''Visualize conv. features as an image (mostly for the 1st layer).
#     Place kernel into a grid, with some paddings between adjacent filters.
#     Args:
#       kernel:            tensor of shape [Y, X, NumChannels, NumKernels]
#       (grid_Y, grid_X):  shape of the grid. Require: NumKernels == grid_Y * grid_X
#                            User is responsible of how to break into two multiples.
#       pad:               number of black pixels around each filter (between them)
# 
#     Return:
#       Tensor of shape [(Y+pad)*grid_Y, (X+pad)*grid_X, NumChannels, 1].
#     '''
#     # pad X and Y
#     x1 = tf.pad(kernel, tf.constant([[pad, 0], [pad, 0], [0, 0], [0, 0]]))
# 
#     # X and Y dimensions, w.r.t. padding
#     Y = kernel.get_shape()[0] + pad
#     X = kernel.get_shape()[1] + pad
#     ch = kernel.get_shape()[2]
# 
#     # put NumKernels to the 1st dimension
#     x2 = tf.transpose(x1, (3, 0, 1, 2))
#     # organize grid on Y axis
#     x3 = tf.reshape(x2, tf.pack([grid_X, Y * grid_Y, X, ch]))
# 
#     # switch X and Y axes
#     x4 = tf.transpose(x3, (0, 2, 1, 3))
#     # organize grid on X axis
#     x5 = tf.reshape(x4, tf.pack([1, X * grid_X, Y * grid_Y, ch]))
# 
#     # back to normal order (not combining with the next step for clarity)
#     x6 = tf.transpose(x5, (2, 1, 3, 0))
# 
#     # to tf.image_summary order [batch_size, height, width, channels],
#     #   where in this case batch_size == 1
#     x7 = tf.transpose(x6, (3, 0, 1, 2))
# 
#     # scale to [0, 1]
#     x_min = tf.reduce_min(x7)
#     x_max = tf.reduce_max(x7)
#     x8 = (x7 - x_min) / (x_max - x_min)
# 
#     return x8


# Path: vae/vaegan_cifar.py
import os
import numpy as np
import tensorflow as tf
from keras import backend as K
from keras.layers import Input, Dense, Reshape, Activation, Convolution2D, Deconvolution2D, LeakyReLU, Flatten, \
    BatchNormalization as BN
from keras.models import Sequential, Model
from keras import initializations
from functools import partial
from utils import cifar
from utils.utils import put_kernels_on_grid



def get_model(sess, image_shape=(32, 32, 3), gf_dim=64, df_dim=64, batch_size=64,
              name="autoencoder"):
    K.set_session(sess)
    with tf.variable_scope(name):
        # sizes
        ch = image_shape[2]
        rows = [4, 8, 16, 32]
        cols = [4, 8, 16, 32]

        # nets
        G = generator(batch_size, gf_dim, ch, rows, cols)
        G.compile("sgd", "mse")
        g_vars = G.trainable_weights
        print "G.shape: ", G.output_shape

        E = encoder(batch_size, df_dim, ch, rows, cols)
        E.compile("sgd", "mse")
        e_vars = E.trainable_weights
        print "E.shape: ", E.output_shape

        D = discriminator(batch_size, df_dim, ch, rows, cols)
        D.compile("sgd", "mse")
        d_vars = D.trainable_weights
        print "D.shape: ", D.output_shape

        Z2 = Input(batch_shape=(batch_size, z_dim), name='more_noise')
        Z = G.input
        Img = D.input

        image_grid = put_kernels_on_grid(tf.transpose(Img, [1, 2, 3, 0]), (8, 8))