import cv2
import os
import numpy as np
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import load_model
import time


class Classifier:
    def __init__(self, is_image, input_path):
        self.is_image = is_image
        self.counter = 0

        self.classifier = load_model('model.h5')
        self.face_detector = cv2.CascadeClassifier(
            'haarcascade_frontalface_default.xml')
        cv2.ocl.setUseOpenCL(False)

        self.labels = {0: "Angry", 1: "Disgusted", 2: "Fearful",
                       3: "Happy", 4: "Neutral", 5: "Sad", 6: "Surprised"}

        self.input_path = input_path
        self.output_path = "result/" + self.input_path.split('/')[-1]
        os.makedirs("result", exist_ok=True)

        print("Models have been loaded")
    # def get_model(self):
    #     model = load_model(model_path, compile=False)
    #     model.compile(optimizer=Adam(learning_rate=1e-3),
    #                 loss='categorical_crossentropy', metrics=['accuracy'])

    def predict(self, frame):

        faces = self.face_detector.detectMultiScale(
            frame, scaleFactor=1.3, minNeighbors=5)

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        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 = self.classifier.predict(cropped_img)
            maxindex = int(np.argmax(prediction))
            cv2.putText(frame, self.labels[maxindex], (x+20, y-60),
                        cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 2, cv2.LINE_AA)

        return frame

    def run(self):
        start_time = time.time()
        if self.is_image:
            image_arr = cv2.imread(self.input_path)
            predicted_frame = self.predict(frame=image_arr)
            cv2.imwrite(self.output_path, predicted_frame)
        else:
            cap = cv2.VideoCapture(self.input_path)

            # Get video properties
            frame_width = int(cap.get(3))
            frame_height = int(cap.get(4))
            fps = int(cap.get(5))

            # Define the codec and create a VideoWriter object to save the output video as .MOV
            fourcc = cv2.VideoWriter_fourcc('a', 'v', 'c', '1')
            out = cv2.VideoWriter(self.output_path, fourcc, fps,
                                  (frame_width, frame_height))

            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                predicted_frame = self.predict(frame=frame)
                out.write(predicted_frame)

        total_time = str(round(time.time() - start_time, 3)) + "s"
        return self.output_path, total_time