#-----AI2.py

import cv2
import random
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing import image
from tensorflow.keras.optimizers import RMSprop
from tensorflow import keras
import tensorflow as tf
import os 
import math
import shutil
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import MaxPooling2D
from tensorflow.keras.layers import Dropout
from tensorflow.keras.utils import to_categorical
from sklearn.model_selection import train_test_split

# 指定亂數種子
seed = 7
np.random.seed(seed)
print('1st cell done...')

# read image
image_size=128

def load_images_from_folder(folder,nfile):
    images = []
    for filename in os.listdir(folder):
        img = cv2.imread(os.path.join(folder,filename), cv2.IMREAD_GRAYSCALE)
        if img is not None:
            img = cv2.resize(img,(image_size,image_size))/255
            images.append(img)
    return np.array(images[:nfile])



# set no. of images for each lable

path='PNG'; ncat=4; nfile=20; 
mod1='model_THUO.h5'
ntest=20; pathT=path+'/TEST/'
LAB=['T','H','U','O']
print('ncat=',ncat,'  nfile=',nfile,' ntest=',ntest)
print('path=',path,' mod1=',mod1,' pathT=',pathT,' LAB=',LAB)

print('load path=',path)
train = load_images_from_folder(path+'/T/',nfile)
print('type(train)=',type(train),train.shape)
Y_train = np.zeros(train.shape[0])
X_train = train
train = load_images_from_folder(path+'/H/',nfile)
Y_train = np.append(Y_train, np.ones(train.shape[0])*1)
X_train = np.append(X_train, train, axis=0)

if(ncat>2):
    train = load_images_from_folder(path+'/U/',nfile)
    Y_train = np.append(Y_train, np.ones(train.shape[0])*2)
    X_train = np.append(X_train, train, axis=0)

if(ncat>3):
    train = load_images_from_folder(path+'/O/',nfile)
    Y_train = np.append(Y_train, np.ones(train.shape[0])*3)
    X_train = np.append(X_train, train, axis=0)

if(ncat>4):
    train = load_images_from_folder(path+'/N5/',nfile)
    Y_train = np.append(Y_train, np.ones(train.shape[0])*4)
    X_train = np.append(X_train, train, axis=0)

if(ncat>5):
    train = load_images_from_folder(path+'/N6/',nfile)
    Y_train = np.append(Y_train, np.ones(train.shape[0])*5)
    X_train = np.append(X_train, train, axis=0)

if(ncat>6):
    train = load_images_from_folder(path+'/N7/',nfile)
    Y_train = np.append(Y_train, np.ones(train.shape[0])*6)
    X_train = np.append(X_train, train, axis=0)

if(ncat>7):
    train = load_images_from_folder(path+'/N8/',nfile)
    Y_train = np.append(Y_train, np.ones(train.shape[0])*7)
    X_train = np.append(X_train, train, axis=0)



print('X_train.shape_5=',X_train.shape); 
print('y_train.shape_5=',Y_train.shape)
#print(Y_train)

# One-hot編碼
Y_train = to_categorical(Y_train)
print('Y_train.shape2=',Y_train.shape)
#print('Y2_train=',Y_train)
#print(Y_train)

# 將圖片轉換成 4D 張量
X_train = X_train.reshape(X_train.shape[0], image_size, image_size, 1).astype("float32")

X_train, X_test, Y_train, Y_test = train_test_split(X_train, Y_train, test_size=0.1, random_state=42)
print('show TEST dat : Fig and Label')

'''
for i in range(0,10):
    img = X_train[i,:,:,0]
    print(i,img.shape)
    cv2.imshow('img',img)
    cv2.waitKey(0)
'''

# 定義模型
model = Sequential()
model.add(Conv2D(16,(3,3),activation = 'relu', input_shape = (image_size, image_size,1)))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Conv2D(32,(3,3),activation = 'relu'))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Flatten())
model.add(Dense(64,activation='relu'))
model.add(Dropout(0.2))
model.add(Dense(32,activation='relu'))
model.add(Dense(ncat,activation='softmax'))

model.summary()

model.compile(loss = 'categorical_crossentropy',
             optimizer = "adam",
             metrics = ['accuracy'])

print(X_train.shape)
print(Y_train.shape)

history = model.fit(X_train,Y_train,shuffle=True, validation_split=0.1,
                                batch_size=128, epochs=60) 

# 顯示訓練和驗證準確度 使用accuracy instead of acc
acc = history.history["accuracy"]
epochs = range(1, len(acc)+1)
val_acc = history.history["val_accuracy"]
'''
plt.plot(epochs, acc, "bo-", label="Training Acc")
plt.plot(epochs, val_acc, "ro--", label="Validation Acc")
plt.title("Training and Validation Accuracy")
plt.xlabel("Epochs")
plt.ylabel("Accuracy")
plt.legend()
plt.show()
'''

'''
# 顯示訓練和驗證損失
loss = history.history["loss"]
epochs = range(1, len(loss)+1)
val_loss = history.history["val_loss"]
plt.plot(epochs, loss, "bo-", label="Training Loss")
plt.plot(epochs, val_loss, "ro--", label="Validation Loss")
plt.title("Training and Validation Loss")
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.legend()
plt.show()
'''


# 評估模型
print(X_train.shape)
print(Y_train.shape)
print(X_test.shape)
print(Y_test.shape)
print("\nTesting ...")
loss, accuracy = model.evaluate(X_train, Y_train, verbose=0)
print("訓練資料集的準確度 = {:.2f}".format(accuracy))
loss, accuracy = model.evaluate(X_test, Y_test, verbose=0)
print("測試資料集的準確度 = {:.2f}".format(accuracy))

model.save(mod1)
print('model ',mod1,' is saved...')

#------- you split this into a new cell if you want a frequent test run-------

model = keras.models.load_model(mod1)
print('load ',mod1,' ... done')

# read images

test = load_images_from_folder(pathT,ntest)
print('load is done path=',pathT)
X_test = test
X_test = X_test.reshape(X_test.shape[0], image_size, image_size, 1).astype("float32")
print('X_test:type,len,shape=',type(X_test),len(X_test),X_test.shape)
pred = np.argmax(model.predict(X_test),axis=1)      
print()
#print(X_test)
print('pred=',pred)
for i in range(ntest):
    img = X_test[i,:,:,0]
    print(i,'pred[i]=',pred[i],' LAB=',LAB[pred[i]])
    cv2.imshow('img',img)
    cv2.waitKey(0)