import numpy as np
import cv2
import torch
from .exp_recognition_model import *
from PIL import Image
import os

#############################################################################################################################
#   Caution: Don't change any of the filenames, function names and definitions                                              #
#############################################################################################################################

current_path = os.path.dirname(os.path.abspath(__file__))


# ============================================================
# DEVICE
# ============================================================

if torch.backends.mps.is_available():
    device = torch.device("mps")
elif torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")


# ============================================================
# FACE DETECTION
# ============================================================

def detected_face(image):

    eye_haar = current_path + '/haarcascade_eye.xml'
    face_haar = current_path + '/haarcascade_frontalface_default.xml'

    face_cascade = cv2.CascadeClassifier(face_haar)
    eye_cascade = cv2.CascadeClassifier(eye_haar)

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=1.3,
        minNeighbors=5
    )

    face_areas = []

    images = []

    required_image = 0

    for i, (x, y, w, h) in enumerate(faces):

        face_cropped = gray[y:y+h, x:x+w]

        face_areas.append(w * h)

        images.append(face_cropped)

        required_image = images[np.argmax(face_areas)]

        required_image = Image.fromarray(required_image)

    return required_image


# ============================================================
# LOAD MODEL
# ============================================================

_model = None


def load_model():

    global _model

    if _model is None:

        model_path = current_path + '/expression_model.pth'
        checkpoint = torch.load(
            model_path,
            map_location=device
        )

        model = facExpRec(
            num_classes=len(classes)
        ).to(device)

        if isinstance(checkpoint, dict):

            if 'model_state_dict' in checkpoint:

                model.load_state_dict(
                    checkpoint['model_state_dict']
                )

            elif 'net_dict' in checkpoint:

                model.load_state_dict(
                    checkpoint['net_dict']
                )

            else:

                model.load_state_dict(checkpoint)

        else:

            model.load_state_dict(checkpoint)

        model.eval()

        _model = model

    return _model


# ============================================================
# EXPRESSION PREDICTION
# ============================================================

def get_expression(img):

    model = load_model()

    face = detected_face(img)

    if face == 0:

        gray = cv2.cvtColor(
            img,
            cv2.COLOR_BGR2GRAY
        )

        face = Image.fromarray(gray)

    image_tensor = trnscm(face)

    image_tensor = image_tensor.unsqueeze(0).to(device)

    with torch.no_grad():

        outputs = model(image_tensor)

        probabilities = torch.softmax(
            outputs,
            dim=1
        )

        confidence, predicted = torch.max(
            probabilities,
            1
        )

    predicted_idx = predicted.item()

    expression = classes[predicted_idx]

    return expression