from keras.preprocessing.image import ImageDataGenerator
datagen = ImageDataGenerator(
    shear_range = 0.2,
    zoom_range = 0.2,
  horizontal_flip = True
)

from keras.layers.convolutional_recurrent import ConvLSTM2D
import tensorflow.keras.optimizers
import keras

from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten, BatchNormalization, MaxPool2D 
from keras.layers import Conv2D, MaxPooling2D
from keras.applications.resnet import ResNet50
from keras.applications.vgg16 import VGG16

def createModel2():
  input_shape = (pic_height, pic_width, 3)
  model = Sequential()

  # 由於我們是圖像，加入 Convolution Layer 必須用 "2D"，
  # kernel_size 用 (3,3) 則是因為我圖像已經相對小了，因此要對細節好好處理
  # padding 對超出邊界的視窗直接補 0 即可
  model.add(Conv2D(32, kernel_size = (3,3), padding = "Same",
            activation="relu", input_shape = input_shape))
  model.add(Conv2D(32, kernel_size=(3,3), padding = "Same", activation="relu"))

  # 在一般情況下 pool_size 設定成 2 就有不錯效果
  model.add(MaxPool2D(pool_size=(2,2)))
  model.add(BatchNormalization())
  model.add(Conv2D(64, kernel_size=(3,3), padding="Same", activation="relu"))
  model.add(Conv2D(64, kernel_size=(3,3), padding="Same", activation="relu"))
  model.add(MaxPool2D(pool_size=(2,2)))

  # 將前面每一層的 layer 的 feature 都做 Normalization，確保不會讓前面的 layer 權重太大，
  # 由於每一層 layer 都會影響到下一層，因此會容易讓前面一層的 layer 影響被放大，就有可能 overfitting 
  #可以參考 [ML筆記] Batch Normalization by 陳雲濤的部落格
  model.add(BatchNormalization())
  model.add(Conv2D(86, kernel_size = (3,3), padding = "Same", activation="relu"))
  model.add(Conv2D(86, kernel_size=(3,3), padding = "Same", activation="relu"))
  model.add(MaxPool2D(pool_size=(2,2)))
  model.add(BatchNormalization())

  # 將多維的輸入轉化成一維
  model.add(Flatten())
  # 將資料變成 1024 個特徵
  model.add(Dense(1024, activation="relu"))
  model.add(Dense(512, activation="relu"))
  
  # 由於我們最後辨識 50 個腳色，因此大小就是 50
  model.add(Dense(50, activation="softmax")) #因為我們資料是機率，所以用 softmax，可以顯現出 50 個 character 機率

  #define the optimizer
  optimizer = "Adam"
  model.compile(optimizer=optimizer, loss="categorical_crossentropy", metrics=["accuracy"])
  return model

from keras.callbacks import ModelCheckpoint, LearningRateScheduler, EarlyStopping

model = createModel2()
model.summary()

lr = 0.001

batch_size = 32
epoch = 50
def lr_schedule(epoch):
  return lr*(0.1**(int(epoch/10)))

# 會將最好的 model 儲存，並且如果 epoch 沒有更好時，就會讓此 model 暫停
history = model.fit(x_train, y_train,
                    batch_size = batch_size,
                    epochs = epoch,
                    validation_data = (x_test, y_test),
                    shuffle = True,
                    callbacks=[LearningRateScheduler(lr_schedule), 
                               ModelCheckpoint("model2.h5", save_best_only=True, period=10),
                               EarlyStopping(monitor="val_accuracy", patience=3, mode= "auto")])

import matplotlib.pyplot as plt

def plot_train_history(history, train_metrics, val_metrics):
  plt.plot(history.history.get(train_metrics), '-o')
  plt.plot(history.history.get(val_metrics, '-o'))
  plt.ylabel(train_metrics)
  plt.xlabel("epoch")
  plt.legend(["train", "validation"])

plt.figure(figsize=(12,4))
plt.subplot(1,2,1)
plot_train_history(history, "loss", "val_loss")

plt.subplot(1,2,2)
plot_train_history(history, "accuracy", "val_accuracy")
plt.show()

import sklearn
predict = model.predict(x_test)
predict = np.argmax(predict, axis=1)
ans = np.argmax(y_test, axis = 1)
print("\n", sklearn.metrics.classification_report(ans, predict, target_names=list(characters.values())), sep='')

from keras.models import load_model
from google.colab.patches import cv2_imshow
def read_images(path):
  images=[]
  for i in range(10791):
  #for i in range(3):
    print(path+str(i+1)+'.jpg')
    image = cv2.resize(cv2.imread(path+str(i+1)+'.jpg'), (pic_height, pic_width))
    #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    #cv2_imshow(image)
    images.append(image)
  images = np.array(images, dtype=np.float32) / 255
  return images

model = load_model("model2.h5")
#imgpath = "/content/drive/My Drive/NTUT/大三上-深度學習/Predict_Simposons/theSimpsons-test/test/" #t108AB0008@ntut.org.tw
imgpath = "/content/drive/.shortcut-targets-by-id/1EmT_yvIKvgl-iNKBU0Wu2XOFjOPsxhyu/大三上-深度學習/Predict_Simposons/theSimpsons-test/test/" #others

test_images = read_images(imgpath)
print(test_images.shape)
predict = model.predict(test_images)
predict = np.argmax(predict, axis=1) #找第一個機率最大的
print(characters[int(predict[2])])

with open("predict2.csv", "w") as f:
  f.write("id,character\n")
  for i in range(len(predict)):
    text = str(i+1) + "," + characters[int(predict[i])] + "\n"
    f.write(text)