import numpy as np

from keras.datasets import mnist

from keras.models import Sequential

from keras.layers import Dense, Dropout, Flatten

from keras.layers.convolutional import Conv2D, MaxPooling2D

(x_train, y_train), (x_test, y_test) = mnist.load_data()

print(x_train[0].shape)

print(y_train)

########################### x 处理 ##################################

# 将训练集合中的数字变成标准的四维张量形式（样本数量、长、宽、深(灰度图 1））

# 并将像素值变成浮点格式

width = 28

height = 28
depth = 1

x_train = x_train.reshape(x_train.shape[0], width, height, depth).astype('float32')

x_test = x_test.reshape(x_test.shape[0], width, height, depth).astype('float32')

# 归一化处理，将像素值控制在 0 - 1

x_train /= 255

x_test /= 255

classes = 10

####################### y 处理 #######################################

# one host 编码

def tran_y(y):

    y_ohe = np.zeros(10)

    y_ohe[y] = 1

    return y_ohe

# 标签将 one-hot 编码重排

y_train_ohe = np.array([tran_y(y_train[i]) for i in range(len(y_train))])

y_test_ohe = np.array([tran_y(y_train[i]) for i in range(len(y_test))])

###################### 搭建卷积神经网络 ###############################

model = Sequential()

# 添加卷积层，构造 64 个过滤器，过滤器范围 3x3x1, 过滤器步长为 1， 图像四周补一圈 0， 并用 relu 非线性变换

model.add(Conv2D(filters=64, kernel_size=(3,3), strides=(1,1), padding='same', input_shape=(width, height, 1), activation='relu'))

# 添加 Max_Pooling , 2 x 2 取最大值

model.add(MaxPooling2D(pool_size=(2, 2)))

# 设立 Dropout ， 将概率设为 0.5

model.add(Dropout(0.5))

#重复构造, 搭建神经网络

model.add(Conv2D(128, kernel_size=(3, 3), strides=(1,1), padding='same', activation='relu'))

model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Dropout(0.5))

model.add(Conv2D(256, kernel_size=(3,3), strides=(1, 1), padding='same', activation='relu'))

model.add((MaxPooling2D(pool_size=(2, 2))))

model.add(Dropout(0.5))

# 将当前节点展平, 构造全连神经网络

model.add(Flatten())

# 构造全连接神经网络

model.add(Dense(128, activation='relu'))

model.add(Dense(64, activation='relu'))

model.add(Dense(32, activation='reul'))

model.add(Dense(classes, activation='softmax'))

################################ 编译模型 ##########################

# 一般，分类问题的损失函数才有交叉熵 (Cross Entropy)

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

######################### 训练模型 ################################

model.fit(x_train, y_train_ohe, validation_data=(x_test, y_test_ohe), epochs=20, batch_size=128)

######################## 评价模型 ################################

scores = model.evaluate(x_test, y_test_ohe, verbose=0)

######################## 保持模型与权重 ################################

# 保持整个模型（包括结构、权重）

model.save("mnist_model.h5")