import numpy as np
from keras.utils import to_categorical
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.layers import Input, Flatten, Dense, Activation, Conv2D, LeakyReLU, Dropout, BatchNormalization, MaxPooling2D
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.optimizers import Adam
from keras import backend as K
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import os

#loss 그래프 그리기
def plot_loss(history):
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('Model loss')
    plt.ylabel('Loss')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'val'], loc=0)

#accuracy 그래프 그리기
def plot_acc(history):
    plt.plot(history.history['accuracy'])
    plt.plot(history.history['val_accuracy'])
    plt.title('Model accuracy')
    plt.ylabel('accuracy')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'val'], loc=0)

img_rows, img_cols = 32, 32

nb_filters = 32


save_dir = os.path.join(os.getcwd(), 'saved_model')
model_name = 'cifar10_model.h5'

#데이터 로드
(x_train, y_train), (x_test, y_test) = cifar10.load_data()

if K.image_dim_ordering() == 'th':
    x_train = x_train.reshape(x_train.shape[0], 3, img_rows, img_cols)
    x_test = x_test.reshape(x_test.shape[0], 3, img_rows, img_cols)
    input_shape = (3, img_rows, img_cols)
else:
    x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 3)
    x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 3)
    input_shape = (img_rows, img_cols, 3)

#샘플 이미지 출력
CLASSES = np.array(['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'])
actual_single = CLASSES[y_train]
plt.imshow(x_train[20], interpolation='bicubic')
tmp = "Label:" + str(actual_single[20])
plt.title(tmp, fontsize=30)
plt.tight_layout()
plt.show()

NUM_CLASSES = 10

x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

y_train = to_categorical(y_train, NUM_CLASSES)
y_test = to_categorical(y_test, NUM_CLASSES)

#모델 만들기
input_layer = Input(shape=input_shape)

x = Conv2D(filters=32, kernel_size=3, strides=1, padding='same')(input_layer)
x = BatchNormalization()(x)
x = LeakyReLU()(x)
x = Conv2D(filters=32, kernel_size=3, strides=2, padding='same')(x)
x = LeakyReLU()(x)
x = Dropout(0.25)(x)
x = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x)
x = LeakyReLU()(x)
x = Dropout(0.5)(x)
x = Conv2D(filters=64, kernel_size=3, strides=2, padding='same')(x)
x = LeakyReLU()(x)
x = Dropout(0.25)(x)
x = MaxPooling2D(pool_size=(2,2))(x)
x = Flatten()(x)
x = Dense(128)(x)
x = LeakyReLU()(x)
x = Dense(NUM_CLASSES)(x)
output_layer = Activation('softmax')(x)
model = Model(input_layer, output_layer)
model.summary()

#모델 컴파일
opt = Adam(lr=0.0005)
model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])

#data aug
datagen = ImageDataGenerator(rotation_range=0.1, width_shift_range=0.1, height_shift_range=0.1, horizontal_flip=True)

#모델 학습(aug적용)
history = model.fit_generator(datagen.flow(x_train, y_train, batch_size=16),
                              steps_per_epoch=x_train.shape[0]/16, epochs=12,
                              validation_data=(x_test, y_test),workers=4)
#모델 훈련
#history = model.fit(x_train, y_train, batch_size=16, epochs=10, verbose=1, validation_split=0.2)

#모델 저장
if not os.path.isdir(save_dir):
    os.mkdir(save_dir)
model_path = os.path.join(save_dir, model_name)
model.save(model_path)

#모델 평가
print('\nAccuracy: {:.4f}'.format(model.evaluate(x_test,y_test, verbose=0)[1]))
print('\nloss: {:.4f}'.format(model.evaluate(x_test,y_test, verbose=0)[0]))

plot_acc(history)
plt.show()
plot_loss(history)
plt.show()

Y_pred = model.predict(x_test)
plt.figure(figsize=(10, 10))
CLASSES = np.array(['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'])
for i in range(25):
    plt.subplot(5, 5, i+1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    plt.imshow(np.reshape(x_test[i], [32,32,3]))
    plt.xlabel("act = " + CLASSES[np.argmax(y_test[i])] + "\npred = " + CLASSES[np.argmax(Y_pred[i])])
    plt.tight_layout()
plt.show()