import os
import numpy as np

from keras.models import Sequential
from keras.optimizers import Adadelta
from keras.losses import categorical_crossentropy
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import Flatten
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.models import model_from_yaml


class FERGModel(object):
    """
    FERG-db Model (https://grail.cs.washington.edu/projects/deepexpr/ferg-db.html)

    - 96*96 pixel grayscale images
    """

    def __init__(self, cfg):
        super().__init__()
        self.config = cfg
        self.model = None
    
    def get_model_structure(self):
        model = Sequential()
        model.add(Conv2D(
            32,
            kernel_size=(3, 3),
            activation='relu',
            input_shape=(self.config['input_dim'][0], self.config['input_dim'][1], 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(Flatten())
        model.add(Dense(
            128,
            activation='relu'
        ))
        model.add(Dropout(0.5))
        model.add(Dense(
            len(self.config['classes']),
            activation='softmax'
        ))
        model.compile(
            optimizer=Adadelta(),
            loss=categorical_crossentropy,
            metrics=['accuracy']
        )

        return model

    def load_model_weights(self):
        self.model = self.get_model_structure()
        self.model.load_weights(os.path.join(os.path.dirname(__file__), self.config['file_name']))

    def get_face_emotion(self, face_image):
        out = np.asarray(face_image.resize((self.config['input_dim'][0], self.config['input_dim'][1])), dtype='float32')
        out /= 255
        out = out.reshape((1, self.config['input_dim'][0], self.config['input_dim'][1], 1))
        predicted_class_index = self.model.predict(out)

        return self.config['classes'][np.argmax(predicted_class_index)]

numpy.argmax