# coding=utf8

import numpy as np
from keras.datasets import mnist
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers import Dense, Activation, Convolution2D, MaxPooling2D, Flatten
from keras.optimizers import Adam
from keras import models

np.random.seed(1337)  # for reproducibility

# 加载数据集，下载的数据保存在'~/.keras/datasets/'
(X_train, y_train), (X_test, y_test) = mnist.load_data()
# print("X_train.shape", X_train.shape, "y_train.shape", y_train.shape)
# print("X_train", X_train[0], "y_train", y_train[0])

# 处理数据（对数据进行归一化）
X_train = X_train.reshape(-1, 28, 28, 1) / 255.  # 为了让激励函数更加有效
X_test = X_test.reshape(-1, 28, 28, 1) / 255.  #
y_train = np_utils.to_categorical(y_train, num_classes=10)
y_test = np_utils.to_categorical(y_test, num_classes=10)
# print("y_train.shape", y_train.shape, "y_train", y_train[0])

# 建立一个model
model = Sequential()

# 添加第一层卷积层
model.add(Convolution2D(
    batch_input_shape=(None, 28, 28, 1),  # 输入源的shape
    filters=32,  # 过滤器的数量(卷积核)
    kernel_size=5,  # 过滤器的大小
    strides=1,  # 过滤器移动的步长
    padding='same',  # Padding的方法
))
print(model.output)

# 给model添加一个激励函数
model.add(Activation('relu'))

# 添加一个(max pooling) output shape (32, 14, 14)
model.add(MaxPooling2D(
    pool_size=2,  # 池化层的大小
    strides=2,  # 池化移动的步长
    padding='same',  # Padding的方法
))
print(model.output)

# 添加第二层卷积层 output shape (64, 14, 14)
model.add(Convolution2D(
    filters=64,
    kernel_size=5,
    strides=1,
    padding='same'
))
print(model.output)

# 给model添加一个激励函数
model.add(Activation('relu'))

# # 添加一个(max pooling) output shape (64, 7, 7)
model.add(MaxPooling2D(
    pool_size=2,
    strides=2,
    padding='same',
))
print(model.output)

# 设置第一个全连接层
model.add(Flatten())  # Flatten层用来将输入“压平”

print(model.output)

model.add(Dense(1024))  # 全连接层
model.add(Activation('relu'))  # 给model添加一个激励函数

# 设置第二个全连接层
model.add(Dense(10))
model.add(Activation('softmax'))

# 定义一个优化器，设置一个学习效率
adam = Adam(lr=1e-4)

# 编译模型指定优化器，损失函数及评价指标
model.compile(optimizer=adam,
              loss='categorical_crossentropy',
              metrics=['accuracy'])

print('Training ------------')
model.fit(X_train, y_train, epochs=5, batch_size=64)

print('\nTesting ------------')
loss, accuracy = model.evaluate(X_test, y_test)

print('\ntest loss: ', loss)
print('\ntest accuracy: ', accuracy)

models.save_model(model, 'mnist_cnn_keras.h5')