import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib as mpl
import os
import numpy as np
import pandas as pd
import pickle
import keras
from keras.models import Sequential
from keras.layers import Conv2D, MaxPool2D, Flatten,BatchNormalization
from keras.layers import Dense, Dropout
from keras.preprocessing.image import ImageDataGenerator
from keras.datasets import cifar10

# 初始化参数
batch_size = 32
epochs = 100  # 迭代次数
num_classes = 10  # 分类数
data_augmentation = False  # 数据增强
model_dir = os.path.join(os.getcwd(), 'saved_models')  # 保存路径
model_name = 'keras_cifar10_trained_model.h5'  # 文件名

# 加载并检查数据集
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
print('x_train.shape = ' + str(x_train.shape))
print('y_train.shape = ' + str(x_test.shape))
print('train samples = ' + str(x_train.shape[0]))
print('test samples = ' + str(x_test.shape[0]))

# 数据预处理
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
x_train = x_train.astype('float32')
x_train /= 255.
x_test = x_test.astype('float32')
x_test /= 255.

# 搭建Sequential模型
model = keras.Sequential()  # conv -> conv -> maxpool -> dropout -> conv -> conv -> maxpool -> dropout -> flatten -> fc -> dropout -> fc -> softmax

model.add(Conv2D(64, (3, 3), padding='same', activation='relu', input_shape=x_train.shape[1:]))
model.add(Conv2D(64, (3, 3), padding="same",activation='relu'))
model.add(BatchNormalization())
model.add(MaxPool2D(pool_size=(2, 2)))
model.add(Conv2D(128, (3, 3),padding="same",activation='relu'))
model.add(BatchNormalization())
model.add(Conv2D(128, (3, 3),padding="same",activation='relu'))
model.add(BatchNormalization())
model.add(MaxPool2D(pool_size=(2, 2)))
model.add(Conv2D(256, (3, 3),padding="same",activation='relu'))
model.add(BatchNormalization())
model.add(Conv2D(256, (3, 3),padding="same",activation='relu'))
model.add(BatchNormalization())
model.add(MaxPool2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(512, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(num_classes, activation='softmax'))

# 模型编译
opt = keras.optimizers.rmsprop(lr=0.0001, decay=1e-6)  # 优化器，这里选择使用
model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])

# 训练模型
if not data_augmentation:
    print("Not using data augementation")
    model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_data=(x_test, y_test),
              shuffle=True)
else:
    print("Using real-time data augmentation")  # 使用实时的数据增加
    data_generate = ImageDataGenerator(featurewise_center=False,  # 将输入数据的均值设置为0
                                       samplewise_center=False,  # 将每个样本的均值设置为0
                                       featurewise_std_normalization=False,  # 将输入除以数据标准差，逐特征进行
                                       samplewise_std_normalization=False,  # 将每个输出除以其标准差
                                       zca_epsilon=1e-6,  # ZCA白化的epsilon值，默认为1e-6
                                       zca_whitening=False,  # 是否应用ZCA白化
                                       rotation_range=0,  # 随机旋转的度数范围，输入为整数
                                       width_shift_range=0.1,  # 左右平移，输入为浮点数，大于1时输出为像素值
                                       height_shift_range=0.1,  # 上下平移，输入为浮点数，大于1时输出为像素值
                                       shear_range=0.,  # 剪切强度，输入为浮点数
                                       zoom_range=0.,  # 随机缩放，输入为浮点数
                                       channel_shift_range=0.,  # 随机通道转换范围，输入为浮点数
                                       fill_mode='nearest',  # 输入边界以外点的填充方式，还有constant,reflect,wrap三种填充方式
                                       cval=0.,  # 用于填充的值，当fill_mode='constant'时生效
                                       horizontal_flip=True,  # 随机水平翻转
                                       vertical_flip=False,  # 随机垂直翻转
                                       rescale=None,  # 重随放因子，为None或0时不进行缩放
                                       preprocessing_function=None,  # 应用于每个输入的函数
                                       data_format=None,  # 图像数据格式，默认为channels_last
                                       validation_split=0.0)

    # 使用实时数据增强的batch对模型进行拟合
    model.fit_generator(data_generate.flow(x_train, y_train, batch_size), steps_per_epoch=len(x_train) // batch_size,
                        epochs=epochs)

# 保存模型
if not os.path.isdir(model_dir):  # 目录是否存在，不存在则创建相应的目录
    os.makedirs(model_dir)
    model_path = os.path.join(model_dir, model_name)
    model.save(model_path)
    print("Saved trained model at: " + str(model_path))

# 模型预测
scores = model.evaluate(x_test, y_test, batch_size=batch_size, verbose=1)  # 以进度条的方式显示迭代过程
print("Test loss is: " + str(scores[0]))
print("Test accuracy is: " + str(scores[1]))