# this is a slightly modified version of the src/web/emotion_gender_processor.py to provide control over paths and other changes (see README.md)

import os, logging, json

import cv2
from keras.models import load_model
import numpy as np

from utils.datasets import get_labels
from utils.inference import detect_faces
from utils.inference import draw_text
from utils.inference import draw_bounding_box
from utils.inference import apply_offsets
from utils.inference import load_detection_model
from utils.preprocessor import preprocess_input

class EGProcessor:
    # hyper-parameters for bounding boxes shape
    gender_offsets = (30, 60)
    gender_offsets = (10, 10)
    emotion_offsets = (20, 40)
    emotion_offsets = (0, 0)

    valid = False

    def __init__(self, dirname, models_dir = 'trained_models'):
        try:
            # parameters for loading data and images
            self.emotion_labels = get_labels('fer2013')
            self.gender_labels = get_labels('imdb')

            # loading models
            self.face_detection = load_detection_model(models_dir + '/detection_models/haarcascade_frontalface_default.xml')
            self.emotion_classifier = load_model(models_dir + '/emotion_models/fer2013_mini_XCEPTION.102-0.66.hdf5', compile=False)
            self.gender_classifier = load_model(models_dir + '/simple_CNN.81-0.96.hdf5', compile=False)

            if not os.path.exists(dirname):
                os.mkdir(dirname)

            self.valid = True
            self.dirname = dirname
        except Exception as err:
            logging.error('Error initializing emotion gender processor: "{0}"'.format(err))
            self.valid = False

    def isValid(self):
        return self.valid

    def processImage(self, image, result_fname = None, type = 'png', detect_emotion = True):
        result = []

        # make sure result is not png if we have no file to save to
        if result_fname == None:
            type = 'mem'

        try:
            # getting input model shapes for inference
            emotion_target_size = self.emotion_classifier.input_shape[1:3]
            gender_target_size = self.gender_classifier.input_shape[1:3]

            # loading images
            image_array = np.fromstring(image, np.uint8)
            unchanged_image = cv2.imdecode(image_array, cv2.IMREAD_UNCHANGED)

            rgb_image = cv2.cvtColor(unchanged_image, cv2.COLOR_BGR2RGB)
            rgb_face = None

            gray_image = cv2.cvtColor(unchanged_image, cv2.COLOR_BGR2GRAY)
            gray_face = None

            faces = detect_faces(self.face_detection, gray_image)
            for face_coordinates in faces:
                x1, x2, y1, y2 = apply_offsets(face_coordinates, self.gender_offsets)
                rgb_face = rgb_image[y1:y2, x1:x2]

                if detect_emotion:
                    x1, x2, y1, y2 = apply_offsets(face_coordinates, self.emotion_offsets)
                    gray_face = gray_image[y1:y2, x1:x2]

                try:
                    rgb_face = cv2.resize(rgb_face, (gender_target_size))

                    if detect_emotion:
                        gray_face = cv2.resize(gray_face, (emotion_target_size))
                except:
                    continue

                rgb_face = np.expand_dims(preprocess_input(rgb_face, False), 0)
                gender_text = self.gender_labels[np.argmax(self.gender_classifier.predict(rgb_face))]

                if gender_text == self.gender_labels[0]:
                    color = (0, 0, 255)
                else:
                    color = (255, 0, 0)

                if type == 'png':
                    draw_bounding_box(face_coordinates, rgb_image, color)
                    draw_text(face_coordinates, rgb_image, gender_text, color, 0, -20, 1, 2)

                detection_result = {
                    "gender": gender_text,
                    "cords": face_coordinates.tolist()
                }

                if detect_emotion:
                    gray_face = np.expand_dims(np.expand_dims(preprocess_input(gray_face, True), 0), -1)
                    detection_result["emotion"] = self.emotion_labels[np.argmax(self.emotion_classifier.predict(gray_face))]

                    if type == 'png':
                        draw_text(face_coordinates, rgb_image, detection_result["emotion"], color, 0, -50, 1, 2)

                result.append(detection_result)
        except Exception as err:
            logging.error('Error in emotion gender processor: "{0}"'.format(err))

        if type == 'mem':
            return result

        try:
            if type == 'png':
                bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
                cv2.imwrite(os.path.join(self.dirname, result_fname + '.png'), bgr_image)

            if type == 'json':
                with open(os.path.join(self.dirname, result_fname + '.json'), 'w', encoding='utf-8') as fp:
                    json.dump(result, fp)
        except Exception as err:
            logging.error('Error saving processor results: "{0}"'.format(err))

        return result
