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


class Classifier:
    def __init__(self):
        self.is_image = None
        self.classifier = None
        self.face_detector = None
        self.labels = {0: "Angry", 1: "Disgusted", 2: "Fearful",
                       3: "Happy", 4: "Neutral", 5: "Surprised", 6: "Sad"}
        self.load_models()
    
    def check_file_type(self, is_image, input_path):
        self.is_image = is_image
        self.input_path = input_path
        self.output_path = "result/" + self.input_path.split('/')[-1]
        os.makedirs("result", exist_ok=True)

    def load_models(self):
        # Load VGG
        yaml_file = open('models/VGG19.yaml', 'r')
        self.classifier = model_from_json(yaml_file.read())
        self.classifier.load_weights("models/VGG19.h5")
        yaml_file.close()
        print("VGG19 has been loaded")

        # Load Face Detector
        self.face_detector = cv2.CascadeClassifier(
        'models/haarcascade_frontalface_default.xml')
        cv2.ocl.setUseOpenCL(False)
        print("Face Detector has been loaded")

    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]
            roi_color = frame[y:y + h, x:x + w]
            # Convert grayscale to RGB
            roi_color_rgb = cv2.cvtColor(roi_color, cv2.COLOR_BGR2RGB)
            cropped_img = np.expand_dims(np.expand_dims(
                cv2.resize(roi_color_rgb, (48, 48)), 0), -1)
            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, 1.4, (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