from keras.callbacks import TensorBoard
from keras.models import Sequential
from keras.optimizers import SGD, Adam
from keras.layers import Dense, Flatten, Dropout
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.models import load_model
import keras
import numpy as np
from keras.applications.imagenet_utils import preprocess_input
from keras import backend as K
from keras.datasets import cifar10
from tensorflow.examples.tutorials.mnist import input_data
K.clear_session()
mnist = input_data.read_data_sets("MNIST_DATA", one_hot=True)
class AlexModel:
    #初始化引數
    def __init__(self, epochs, batch_size):
        """
        :param epochs: 訓練集迭代的輪數
        :param batch_size: 每次訓練的樣本的個數
        """
        self.epochs = epochs
        self.batch_size = batch_size
        # 儲存訓練過程中的精度和誤差
        self.train_accuracy_and_loss = None
    # 建立模型
    def build_model(self):
        """
        建立模型, 基於alexnet
        :return: 
        """

        model = Sequential()
        #第一層卷積網路，使用96個卷積核，大小為11x11步長為4， 要求輸入 1個通道，啟用函式使用relu
        model.add(Conv2D(96, (11, 11), strides=(4, 4), input_shape=(28, 28, 1), padding='valid', activation='relu',
                         kernel_initializer='uniform'))
        # 池化層
        model.add(MaxPooling2D(pool_size=(3, 3), strides=(2, 2), padding='valid'))

        # 第二層加邊使用256個5x5的卷積核，加邊，啟用函式為relu
        model.add(Conv2D(256, (5, 5), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
        #使用池化層，步長為2
        model.add(MaxPooling2D(pool_size=(3, 3), strides=(2, 2), padding='same'))

        # 第三層卷積，大小為3x3的卷積核使用384個
        model.add(Conv2D(384, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
        # 第四層卷積,同第三層
        model.add(Conv2D(384, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
        # 第五層卷積使用的卷積核為256個，其他同上
        model.add(Conv2D(256, (3, 3), strides=(1, 1), padding='same', activation='relu', kernel_initializer='uniform'))
        model.add(MaxPooling2D(pool_size=(3, 3), strides=(2, 2), padding='same'))
        # 將卷積展開為神經元
        model.add(Flatten())
        # 第1層隱藏全連線層使用4096個神經元
        model.add(Dense(4096, activation='relu'))
        # dropout正則化
        model.add(Dropout(0.5))
        # 第2層隱藏使用4096個神經元
        model.add(Dense(4096, activation='relu'))
        model.add(Dropout(0.5))
        # 輸出層輸出類別個數
        model.add(Dense(10, activation='softmax'))
        # 選用adam優化器，學習率為0.0003
        adam = Adam(lr=0.0003, decay=1e-6)
        # 編譯模型
        model.compile(loss='categorical_crossentropy', optimizer=adam, metrics=['accuracy'])
        return model
    # 儲存模型
    def save_model_after_train(self):
        model = self.build_model()
        x_train, y_train = mnist.train.images, mnist.train.labels
        x_train = x_train.reshape(55000, 28, 28, 1)
        self.train_accuracy_and_loss = model.fit(x_train, y_train, batch_size=self.batch_size, epochs=self.epochs)
        model.save("model.h5")
    # 載入模型
    def load_model(self):
        return load_model("model.h5")
    # 訓練模型
    def train(self, mnist):
        modle = self.build_model()

        x_train, y_train = mnist.train.images, mnist.train.labels
        x_train = x_train.reshape(55000, 28, 28, 1)
        # {'acc': [], 'loss': []}
        self.train_accuracy_and_loss = modle.fit(x_train, y_train, batch_size=self.batch_size,
                                                 epochs=self.epochs,
                                                 callbacks=[TensorBoard(log_dir='mytensorboard/3')])

        # 獲取訓練過程中的損失（每個epoch）
    def get_train_loss(self):
        return self.train_accuracy_and_loss.history["loss"]
    # 獲取訓練過程中的精度(每個epoch)
    def get_train_accuracy(self):
        return self.train_accuracy_and_loss.history["acc"]
    # 測試集的精度和誤差
    def test_accuracy_and_loss(self):
        """"將訓練好的模型直接拿過來用"
        :return: 返回精度和損失
        """
        model = self.load_model()
        x_test, y_test= mnist.test.images, mnist.test.labels
        x_test = x_test.reshape(10000, 28, 28, 1)
        score = model.evaluate(x_test, y_test, batch_size=32)
        return score[1], score[0]

    
model = AlexModel(epochs=2, batch_size=256)
model.train(mnist)
loss = model.get_train_loss()
acc = model.get_train_accuracy()
print(acc)