# 导入必要的库
import tensorflow as tf
from keras import datasets, layers, models
import matplotlib.pyplot as plt

import matplotlib
matplotlib.use('TkAgg')  # 使用 TkAgg 后端
import matplotlib.pyplot as plt

# 主函数入口
if __name__ == '__main__':
    # 加载 MNIST 数据集。分为训练集和测试集
    (train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()

    # 对图像数据进行标准化处理，将像素值缩放到 0-1 之间
    train_images, test_images = train_images / 255.0, test_images / 255.0
    
    # 打印数据维度信息
    print("data_shpe:",train_images.shape,test_images.shape,train_labels.shape,test_labels.shape)
    """
    输出：((60000, 28, 28), (10000, 28, 28), (60000,), (10000,))
    """

    # 绘制训练图像的示例
    plt.figure(figsize=(20, 10))
    for i in range(20):
        plt.subplot(2, 10, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.imshow(train_images[i], cmap=plt.cm.binary)
        plt.xlabel(train_labels[i])
    plt.show()

    # 调整图像数据的维度，以符合模型的输入要求
    train_images = train_images.reshape((60000, 28, 28, 1))
    test_images = test_images.reshape((10000, 28, 28, 1))

    # 再次打印调整后的数据维度信息
    print("data_shpe:",train_images.shape, test_images.shape, train_labels.shape, test_labels.shape)
    """
    输出：((60000, 28, 28, 1), (10000, 28, 28, 1), (60000,), (10000,))
    """

    # 构建卷积神经网络模型
    model = models.Sequential([
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),

        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(10)
    ])
    model.summary()  # 打印模型结构

    # 编译模型，设置优化器、损失函数和评价指标
    model.compile(
        optimizer='adam',
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=['accuracy'])

    # 训练模型
    history = model.fit(
        train_images,
        train_labels,
        epochs=10,
        )