import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
from tensorflow import keras
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D
import tensorflow.keras.losses as losses
import tensorflow.keras.optimizers as optimizers
from tensorflow.keras import backend as K


def show_train_history(train_history, str_test_acc, train_acc='acc', test_acc='val_acc'):
    # draw with training history
    plt.plot(train_history.history[train_acc])
    plt.plot(train_history.history[test_acc])
    # style the diagram
    plt.title('Train History')
    plt.ylabel('Accuracy')
    plt.xlabel('Epoch')
    plt.legend(['train', 'test'], loc='upper left')
    # save the diagram
    plt.savefig(str_test_acc + '.png')


def enumerate_files(sets, path='All_gray_1_32_32', n_poses=3, n_samples=20):
    filenames, labels = [], []
    # get the set name from sets list
    for name_set in sets:
        # get which pose is now
        for i_pose in range(n_poses):
            # for use with pose index to control the name for 9 separate directories
            for j in range(3):
                dir_name = path + '/' + name_set + \
                    '/000' + str(i_pose * 3 + j) + '/'
                # get the index of sample
                for i_sample in range(n_samples):
                    d = dir_name + '%04d/' % i_sample
                    # list all files in the 'd' directory
                    for f in os.listdir(d):
                        # find the file which the filename extension is 'jpg' type
                        if f.endswith('jpg'):
                            # append the whole path into the filename list
                            filenames.append(d + f)
                            #  append the label number into the labels list
                            labels.append(i_pose)
    return filenames, labels


def read_images(files):
    imgs = []
    # get file in the files list
    for f in files:
        # read the image by the gray-level
        img = cv2.imread(f, cv2.IMREAD_GRAYSCALE)
        # limit range in 0-255, and convert type to float
        img = img/255.0
        # append the image into the 'imgs' list
        imgs.append(img)
    return imgs


def load_datasets(datasets):
    # get all file paths and labels
    files, labels = enumerate_files(datasets)
    # start read all image
    list_of_arrays = read_images(files)
    # convert the list, saving all image, to numpy array
    return np.array(list_of_arrays), labels


def main():
    train_sets = ['Set1', 'Set2', 'Set3']
    test_sets = ['Set4', 'Set5']

    batch_size = 360
    num_classes = 3
    epochs = 200
    input_shape = (32, 32, 1)

    # load all data and labels of training dataset and testing dataset
    x_train, y_train = load_datasets(train_sets)
    x_test, y_test = load_datasets(test_sets)

    # reshape data of training and testing dataset
    x_train = x_train.reshape((x_train.shape[0], 32, 32, 1))
    x_test = x_test.reshape((x_test.shape[0], 32, 32, 1))

    #
    y_train_onehot = to_categorical(y_train, num_classes)
    y_test_onehot = to_categorical(y_test, num_classes)

    # building model
    model = Sequential()
    model.add(Conv2D(32, kernel_size=(3, 3), padding='same',
                     activation='tanh', input_shape=input_shape))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.35))

    model.add(Conv2D(32, kernel_size=(3, 3), padding='same', activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.4))

    model.add(Conv2D(28, kernel_size=(3, 3), padding='same', activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.45))

    model.add(Flatten())
    model.add(Dense(128, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(num_classes, activation='softmax'))

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

    train_history = model.fit(x_train, y_train_onehot, batch_size=batch_size, epochs=epochs,
                              verbose=2, validation_data=(x_test, y_test_onehot))

    score = model.evaluate(x_test, y_test_onehot, verbose=0)
    print(f'Test loss: {score[0]}\nTest accuracy: {score[1]}')

    show_train_history(train_history, str(score[1]))


if __name__ == '__main__':
    main()