# Keras versions
import numpy
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten, Convolution2D, MaxPooling2D
from keras.utils import np_utils
from keras.datasets import mnist

def normalize_data(x, y, convolutional=False):
  a,b,c = x.shape
  if convolutional:
    x = x.reshape(a, b, c, 1).astype('float32')
  else:
    x = x.reshape(a, b*c).astype('float32')
  x /= 255
  y = np_utils.to_categorical(y, 10)
  return x, y

def regular_model(model=Sequential()):
  model.add(Dense(800, input_dim=784, init='normal', activation='relu'))
  model.add(Dense(10, init='normal', activation='softmax'))

  model.compile(loss='categorical_crossentropy', optimizer='SGD', metrics=['accuracy'])
  return model

def convolutional_model(model=Sequential(),
                        kernel_size=(3,3),
                        pool_size=(2,2),
                        img_size=(28, 28),
                        filters_count=32):
  model.add(Convolution2D(filters_count, *kernel_size, border_mode='valid',
   input_shape=(img_size[0], img_size[1], 1)))
  model.add(Activation('relu'))
  model.add(Convolution2D(filters_count, *kernel_size))
  model.add(Activation('relu'))
  model.add(MaxPooling2D(pool_size=pool_size))
  model.add(Dropout(0.25))

  model.add(Flatten())
  model.add(Dense(128))
  model.add(Activation('relu'))
  model.add(Dropout(0.5))
  model.add(Dense(10))
  model.add(Activation('softmax'))

  model.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy'])
  return model

def training_model(x_train, y_train, x_test, y_test, convolutional=False):
  x_train, y_train = normalize_data(x_train, y_train, convolutional=convolutional)
  model = convolutional_model() if convolutional else regular_model()
  model.fit(x_train, y_train, batch_size=32, nb_epoch=25, validation_split=0.1, verbose=True)
  score = model.evaluate(*normalize_data(x_test, y_test, convolutional=convolutional), verbose=1)
  print('Test score:', score[0])
  print('Test accuracy:', score[1])

if __name__ == '__main__':
  (x_train, y_train), (x_test, y_test) = mnist.load_data()
  training_model(x_train, y_train, x_test, y_test, convolutional=True)