import numpy as np
import argparse
import matplotlib.pyplot as plt
import cv2
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.layers import MaxPooling2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

# command line argument
ap = argparse.ArgumentParser()
ap.add_argument("--mode",help="train/display")
mode = ap.parse_args().mode

def combine_gen(*gens):
    while True:
        for g in gens:
            yield next(g)

# plots accuracy and loss curves
def plot_model_history(model_history):
    """
    Plot Accuracy and Loss curves given the model_history
    """
    fig, axs = plt.subplots(1,2,figsize=(15,5))
    # summarize history for accuracy
    axs[0].plot(range(1,len(model_history.history['accuracy'])+1),model_history.history['accuracy'])
    axs[0].plot(range(1,len(model_history.history['val_accuracy'])+1),model_history.history['val_accuracy'])
    axs[0].set_title('Model Accuracy')
    axs[0].set_ylabel('Accuracy')
    axs[0].set_xlabel('Epoch')
    axs[0].set_xticks(np.arange(1,len(model_history.history['accuracy'])+1),len(model_history.history['accuracy'])/10)
    axs[0].legend(['train', 'val'], loc='best')
    # summarize history for loss
    axs[1].plot(range(1,len(model_history.history['loss'])+1),model_history.history['loss'])
    axs[1].plot(range(1,len(model_history.history['val_loss'])+1),model_history.history['val_loss'])
    axs[1].set_title('Model Loss')
    axs[1].set_ylabel('Loss')
    axs[1].set_xlabel('Epoch')
    axs[1].set_xticks(np.arange(1,len(model_history.history['loss'])+1),len(model_history.history['loss'])/10)
    axs[1].legend(['train', 'val'], loc='best')
    fig.savefig('plot.png')
    plt.show()

# Define data generators
train_dir = 'data/train'
val_dir = 'data/test'
eval_dir = 'data/evaluate'
clone_time = 30
num_train = 28709 * clone_time
num_val = 7178
batch_size = 64
num_epoch = 10

train_datagen = ImageDataGenerator(
    rescale=1./255, 
    brightness_range=[0.2,1.5],
    horizontal_flip=True)
val_datagen = ImageDataGenerator(rescale=1./255)
eval_datagen = ImageDataGenerator(rescale=1./255)

train_generators = []
for x in range(clone_time):
    train_generators.append(train_datagen.flow_from_directory(
            train_dir,
            target_size=(48,48),
            batch_size=batch_size,
            color_mode="grayscale",
            class_mode='categorical'))

validation_generator = val_datagen.flow_from_directory(
        val_dir,
        target_size=(48,48),
        batch_size=batch_size,
        color_mode="grayscale",
        class_mode='categorical')

evaluation_generator = eval_datagen.flow_from_directory(
        eval_dir,
        target_size=(48,48),
        batch_size=batch_size,
        color_mode="grayscale",
        class_mode='categorical')


# Create the model
model = Sequential()

model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(48,48,1)))
model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.25))

model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.25))

model.add(Flatten())
model.add(Dense(1024, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(7, activation='softmax'))
model.summary()

# If you want to train the same model or try other models, go for this
if mode == "train":
    model.compile(loss='categorical_crossentropy',optimizer=Adam(lr=0.0001, decay=1e-6),metrics=['accuracy'])
    model_info = model.fit(
            combine_gen(*train_generators),
            steps_per_epoch=num_train // batch_size,
            epochs=num_epoch,
            validation_data=validation_generator,
            validation_steps=num_val // batch_size)
    # plot_model_history(model_info)
    model.save_weights('model-epoch-augmentated{}.h5'.format(num_epoch))

# emotions will be displayed on your face from the webcam feed
elif mode == "display":
    model.load_weights('model-epoch-augmentated{}.h5'.format(num_epoch))

    # prevents openCL usage and unnecessary logging messages
    cv2.ocl.setUseOpenCL(False)

    # dictionary which assigns each label an emotion (alphabetical order)
    emotion_dict = {0: "Angry", 1: "Disgusted", 2: "Fearful", 3: "Happy", 4: "Neutral", 5: "Sad", 6: "Surprised"}

    # start the webcam feed
    cap = cv2.VideoCapture(0)
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    while True:
        # Find haar cascade to draw bounding box around face
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.flip(frame, 1)
        facecasc = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = facecasc.detectMultiScale(gray,scaleFactor=1.3, minNeighbors=5)

        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x, y-50), (x+w, y+h+10), (255, 0, 0), 2)
            roi_gray = gray[y:y + h, x:x + w]
            cropped_img = np.expand_dims(np.expand_dims(cv2.resize(roi_gray, (48, 48)), -1), 0)
            prediction = model.predict(cropped_img)
            np.set_printoptions(formatter={'float': '{: 0.3f}'.format})
            print(prediction)
            # finalArr.append(prediction)
            maxindex = int(np.argmax(prediction))
            cv2.putText(frame, emotion_dict[maxindex], (x+20, y-60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)

        cv2.imshow('Video', cv2.resize(frame,(1600,960),interpolation = cv2.INTER_CUBIC))
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    # np.set_printoptions(formatter={'float': '{: 0.3f}'.format})
    # for x in finalArr:
    #     print(x)
    cv2.destroyAllWindows()
elif mode == "evaluate":
    model.compile(loss='categorical_crossentropy',optimizer=Adam(lr=0.0001, decay=1e-6),metrics=['accuracy'])
    model.load_weights('model-epoch-augmentated{}.h5'.format(num_epoch))
    result = model.evaluate(evaluation_generator)

tensorflow.keras.preprocessing.image.ImageDataGenerator tensorflow.keras.layers.Dense numpy.set_printoptions tensorflow.keras.layers.Conv2D matplotlib.pyplot.subplots tensorflow.keras.layers.MaxPooling2D tensorflow.keras.optimizers.Adam numpy.argmax tensorflow.keras.layers.Dropout matplotlib.pyplot.show tensorflow.keras.models.Sequential tensorflow.keras.layers.Flatten
[(16, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), False, 'import argparse\n'), (60, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'ImageDataGenerator', ([], {'rescale': '(1.0 / 255)', 'brightness_range': '[0.2, 1.5]', 'horizontal_flip': '(True)'}), False, 'from tensorflow.keras.preprocessing.image import ImageDataGenerator\n'), (64, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'ImageDataGenerator', ([], {'rescale': '(1.0 / 255)'}), False, 'from tensorflow.keras.preprocessing.image import ImageDataGenerator\n'), (65, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'ImageDataGenerator', ([], {'rescale': '(1.0 / 255)'}), False, 'from tensorflow.keras.preprocessing.image import ImageDataGenerator\n'), (92, 'tensorflow.keras.models.Sequential', 'Sequential', ([], {}), False, 'from tensorflow.keras.models import Sequential\n'), (30, 'matplotlib.pyplot.subplots', 'plt.subplots', (['(1)', '(2)'], {'figsize': '(15, 5)'}), True, 'import matplotlib.pyplot as plt\n'), (48, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (94, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(32)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""', 'input_shape': '(48, 48, 1)'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (95, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (96, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', ([], {'pool_size': '(2, 2)'}), False, 'from tensorflow.keras.layers import MaxPooling2D\n'), (97, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.25)'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (99, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(128)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (100, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', ([], {'pool_size': '(2, 2)'}), False, 'from tensorflow.keras.layers import MaxPooling2D\n'), (101, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(128)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (102, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', ([], {'pool_size': '(2, 2)'}), False, 'from tensorflow.keras.layers import MaxPooling2D\n'), (103, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.25)'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (105, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (106, 'tensorflow.keras.layers.Dense', 'Dense', (['(1024)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (107, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (108, 'tensorflow.keras.layers.Dense', 'Dense', (['(7)'], {'activation': '"""softmax"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (128, 'cv2.ocl.setUseOpenCL', 'cv2.ocl.setUseOpenCL', (['(False)'], {}), False, 'import cv2\n'), (134, 'cv2.VideoCapture', 'cv2.VideoCapture', (['(0)'], {}), False, 'import cv2\n'), (135, 'cv2.VideoWriter_fourcc', 'cv2.VideoWriter_fourcc', (["*'XVID'"], {}), False, 'import cv2\n'), (165, 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ([], {}), False, 'import cv2\n'), (113, 'tensorflow.keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0001)', 'decay': '(1e-06)'}), False, 'from tensorflow.keras.optimizers import Adam\n'), (141, 'cv2.flip', 'cv2.flip', (['frame', '(1)'], {}), False, 'import cv2\n'), (142, 'cv2.CascadeClassifier', 'cv2.CascadeClassifier', (['"""haarcascade_frontalface_default.xml"""'], {}), False, 'import cv2\n'), (143, 'cv2.cvtColor', 'cv2.cvtColor', (['frame', 'cv2.COLOR_BGR2GRAY'], {}), False, 'import cv2\n'), (147, 'cv2.rectangle', 'cv2.rectangle', (['frame', '(x, y - 50)', '(x + w, y + h + 10)', '(255, 0, 0)', '(2)'], {}), False, 'import cv2\n'), (151, 'numpy.set_printoptions', 'np.set_printoptions', ([], {'formatter': "{'float': '{: 0.3f}'.format}"}), True, 'import numpy as np\n'), (155, 'cv2.putText', 'cv2.putText', (['frame', 'emotion_dict[maxindex]', '(x + 20, y - 60)', 'cv2.FONT_HERSHEY_SIMPLEX', '(1)', '(255, 255, 255)', '(2)', 'cv2.LINE_AA'], {}), False, 'import cv2\n'), (157, 'cv2.resize', 'cv2.resize', (['frame', '(1600, 960)'], {'interpolation': 'cv2.INTER_CUBIC'}), False, 'import cv2\n'), (154, 'numpy.argmax', 'np.argmax', (['prediction'], {}), True, 'import numpy as np\n'), (158, 'cv2.waitKey', 'cv2.waitKey', (['(1)'], {}), False, 'import cv2\n'), (167, 'tensorflow.keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0001)', 'decay': '(1e-06)'}), False, 'from tensorflow.keras.optimizers import Adam\n'), (149, 'cv2.resize', 'cv2.resize', (['roi_gray', '(48, 48)'], {}), False, 'import cv2\n')]