from __future__ import print_function
import matplotlib.pyplot as plt
import numpy as np
import keras
from keras.datasets import mnist
from keras.models import Sequential, load_model
from keras.layers import Flatten, Dense, Dropout
from keras.layers import MaxPooling2D, ZeroPadding2D
from keras.layers.convolutional import Convolution2D
from keras.layers.core import Activation
from keras import backend as K
from keras.utils import np_utils


# input image dimensions
img_x, img_y = 28, 28

# load the MNIST data set, which already splits into train and test sets for us
(x_train, y_train), (x_test, y_test) = mnist.load_data()

print(x_train[0])
plt.imshow(x_train[0],cmap='gray')
plt.show()
# # reshape the data into a 4D tensor - (sample_number, x_img_size, y_img_size, num_channels)
# # because the MNIST is greyscale, we only have a single channel - RGB colour images would have 3
x_train = x_train.reshape(x_train.shape[0], img_x, img_y, 1)
x_test = x_test.reshape(x_test.shape[0], img_x, img_y, 1)
# input_shape = (img_x, img_y, 1)

# # convert the data to the right type
# x_train = x_train.astype('float32')
# x_test = x_test.astype('float32')
# x_train /= 255
# x_test /= 255
# print('x_train shape:', x_train.shape)
# print(x_train.shape[0], 'train samples')
# print(x_test.shape[0], 'test samples')

# model = load_model('cnn_arch_1')
# print(model)


# # Get the first training data
# # Any random image would work as well
# img_to_visualize = x_train[0]
# # Make the image 4d
# img_to_visualize = np.expand_dims(img_to_visualize, axis=0)

# # Function to visualize a layer
# def layer_to_visualize(layer):
    # inputs = model.inputs

    # _convout1_f = K.function(inputs, [layer.output])
    # def convout1_f(X):
        # # The [0] is to disable the training phase flag
        # return _convout1_f([X])

    # convolutions = convout1_f(img_to_visualize)
    # convolutions = np.squeeze(convolutions)

    # print ('Shape of conv:', convolutions.shape)

    # n = convolutions.shape[2]
    # n = int(np.ceil(np.sqrt(n)))

    # # Visualization of each filter of the layer
    # fig = plt.figure(figsize=(12,8))
    # for i in range(len(convolutions[0,0,:])):
        # ax = fig.add_subplot(n,n,i+1)
        # ax.imshow(convolutions[:,:,i], cmap='gray')
    # plt.show()


# print(len(model.layers[:-3]))
# for layer in model.layers[:-3]:
    # layer_to_visualize(layer)
