import openvino as ov
import cv2
import numpy as np

class EmotionModel:
    def __init__(self):
        self.face_compiled_model, self.face_input_layer, self.face_output_layer = self.load_model('face-detection-adas-0001')
        self.emotion_compiled_model, self.emotion_input_layer, self.emotion_output_layer = self.load_model('emotions-recognition-retail-0003')
        

    def load_model(self, model_name):
        model_path = "models/" + model_name + ".xml"
        core = ov.Core()
        model = core.read_model(model=model_path)
        compiled_model = core.compile_model(model=model, device_name="CPU")
        input_layer = compiled_model.input(0)
        output_layer = compiled_model.output(0)
        return compiled_model, input_layer, output_layer

    def preprocess(self, image, input_layer):
 
        input_h, input_w = input_layer.shape[2], input_layer.shape[3]
        input_image = cv2.resize(image, (input_w,input_h))
        input_image = input_image.transpose(2, 0, 1)
        input_image = np.expand_dims(input_image, 0)
        
        return input_image


    def post_process_face(self, result_face, img, conf=0.5):
        boxes = []
        img.copy()                     # 원본 이미지를 수정하는 대신 복사합니다
        h,w,_ = img.shape
        predictions = result_face[0][0]       # 하위 집합 데이터 프레임
        confidence = predictions[:,2]    # conf 값 가져오기 [image_id, label, conf, x_min, y_min, x_max, y_max]
      
        top_predictions = predictions[(confidence>conf)]         # 임계값보다 큰 conf 값을 가진 예측만 선택
        for detection in top_predictions:
            box = (detection[3:7]* np.array([w, h, w, h])).astype("int") # 상자 위치 결정
            box = [0 if i < 0 else i for i in box]
            (xmin, ymin, xmax, ymax) = box   # xmin, ymin, xmax, ymax에 상자 위치 값 지정
            boxes.append(box)
            cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)       # 사각형 만들기
        
        return boxes

    def post_process_emotion(self, result_emotion, img, face_position):
        
        emotions = {
            0:"neutral",
            1:"happy",
            2:"sad",
            3:"surprise",
            4:"anger"
        }
                    
        predictions = result_emotion[0,:,0,0]   
        topresult_index = np.argmax(predictions)
        emotion = emotions[topresult_index]

        font_size = img.shape[0]/1000
        font_thickness = int(img.shape[0]/500)
        text_offset = int(img.shape[0]/30)
       
        cv2.putText(img, emotion,
                    (face_position[0],face_position[1]+text_offset),
                    cv2.FONT_HERSHEY_SIMPLEX, font_size, 
                    (255, 255,255), font_thickness)
        
        return emotion

    def process(self, img):

        input_img = self.preprocess(img, self.face_input_layer)
        result_face = self.face_compiled_model([input_img])[self.face_output_layer]
        boxes = self.post_process_face(result_face, img, conf=0.5)

        if boxes is not None:

            for box in boxes:
                xmin, ymin, xmax, ymax = box
                emotion_input = img[ymin:ymax,xmin:xmax]
                input_img = self.preprocess(emotion_input, self.emotion_input_layer)       
                result_emotion = self.emotion_compiled_model([input_img])[self.emotion_output_layer]
                self.post_process_emotion(result_emotion, img, box)

        return img