"""
Pytests to check facial expression recognition functionality
"""

import os
import pathlib
from typing import List

import cv2
import numpy as np
import pytest
import torch
from facenet_pytorch import MTCNN

from emotiefflib.facial_analysis import EmotiEffLibRecognizer, get_model_list

FILE_DIR = pathlib.Path(__file__).parent.resolve()


def recognize_faces(frame: np.ndarray, device: str) -> List[np.array]:
    """
    Detects faces in the given image and returns the facial images cropped from the original.

    This function reads an image from the specified path, detects faces using the MTCNN
    face detection model, and returns a list of cropped face images.

    Args:
        frame (numpy.ndarray): The image frame in which faces need to be detected.
        device (str): The device to run the MTCNN face detection model on, e.g., 'cpu' or 'cuda'.

    Returns:
        list: A list of numpy arrays, representing a cropped face image from the original image.

    Example:
        faces = recognize_faces('image.jpg', 'cuda')
        # faces contains the cropped face images detected in 'image.jpg'.
    """

    def detect_face(frame: np.ndarray):
        # pylint: disable=unbalanced-tuple-unpacking
        mtcnn = MTCNN(keep_all=False, post_process=False, min_face_size=40, device=device)
        bounding_boxes, probs = mtcnn.detect(frame, landmarks=False)
        if probs[0] is None:
            return []
        bounding_boxes = bounding_boxes[probs > 0.9]
        return bounding_boxes

    bounding_boxes = detect_face(frame)
    facial_images = []
    for bbox in bounding_boxes:
        box = bbox.astype(int)
        x1, y1, x2, y2 = box[0:4]
        facial_images.append(frame[y1:y2, x1:x2, :])
    return facial_images


@pytest.mark.parametrize("model_name", get_model_list())
@pytest.mark.parametrize("engine", ["torch", "onnx"])
def test_one_image_prediction(model_name, engine):
    """
    Simple test with one image
    """
    if model_name == "enet_b0_8_va_mtl" or (
        engine == "onnx" and model_name == "enet_b0_8_best_afew"
    ):
        exp_emotions = ["Happiness", "Anger", "Happiness"]
    else:
        exp_emotions = ["Happiness", "Anger", "Fear"]
    input_file = os.path.join(FILE_DIR, "..", "tests", "test_images", "20180720_174416.jpg")
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    frame_bgr = cv2.imread(input_file)
    frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

    facial_images = recognize_faces(frame, device)

    fer = EmotiEffLibRecognizer(engine=engine, model_name=model_name, device=device)

    emotions = []
    for face_img in facial_images:
        emotion, _ = fer.predict_emotions(face_img, logits=True)
        emotions.append(emotion[0])

    assert emotions == exp_emotions


@pytest.mark.parametrize("model_name", get_model_list())
@pytest.mark.parametrize("engine", ["torch", "onnx"])
def test_one_image_multi_prediction(model_name, engine):
    """
    Simple test with one image and predict_multi_emotions API
    """
    if model_name == "enet_b0_8_va_mtl" or (
        engine == "onnx" and model_name == "enet_b0_8_best_afew"
    ):
        exp_emotions = ["Happiness", "Anger", "Happiness"]
    else:
        exp_emotions = ["Happiness", "Anger", "Fear"]
    input_file = os.path.join(FILE_DIR, "..", "tests", "test_images", "20180720_174416.jpg")
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    frame_bgr = cv2.imread(input_file)
    frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

    facial_images = recognize_faces(frame, device)

    fer = EmotiEffLibRecognizer(engine=engine, model_name=model_name, device=device)

    emotions, _ = fer.predict_emotions(facial_images, logits=True)

    assert emotions == exp_emotions


@pytest.mark.parametrize("model_name", get_model_list())
def test_one_image_features(model_name):
    """
    Compare feature vectors for ONNX and Torch implementations
    """
    input_file = os.path.join(FILE_DIR, "..", "tests", "test_images", "20180720_174416.jpg")
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    frame_bgr = cv2.imread(input_file)
    frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

    facial_images = recognize_faces(frame, device)

    fer_onnx = EmotiEffLibRecognizer(engine="onnx", model_name=model_name, device=device)
    fer_torch = EmotiEffLibRecognizer(engine="torch", model_name=model_name, device=device)

    for face_img in facial_images:
        features_onnx = fer_onnx.extract_features(face_img)
        features_torch = fer_torch.extract_features(face_img)
        assert len(features_onnx) == len(features_torch)


@pytest.mark.parametrize("model_name", get_model_list())
def test_one_image_multi_features(model_name):
    """
    Compare feature vectors for ONNX and Torch implementations with extract_multi_features API
    """
    input_file = os.path.join(FILE_DIR, "..", "tests", "test_images", "20180720_174416.jpg")
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    frame_bgr = cv2.imread(input_file)
    frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

    facial_images = recognize_faces(frame, device)

    fer_onnx = EmotiEffLibRecognizer(engine="onnx", model_name=model_name, device=device)
    fer_torch = EmotiEffLibRecognizer(engine="torch", model_name=model_name, device=device)

    features_onnx = fer_onnx.extract_features(facial_images)
    features_torch = fer_torch.extract_features(facial_images)
    assert features_onnx.shape[0] == 3
    assert features_onnx.shape == features_torch.shape


@pytest.mark.parametrize("model_name", get_model_list())
@pytest.mark.parametrize("engine", ["torch", "onnx"])
def test_inference_affect_net(model_name, engine):
    """
    Run accuracy testing on AffectNet dataset
    """
    files_limit = 100
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"
    fer = EmotiEffLibRecognizer(engine=engine, model_name=model_name, device=device)
    inputs_dir = os.path.join(FILE_DIR, "data", "AffectNet_val")
    input_files = []
    input_labels = []

    for label in os.listdir(inputs_dir):
        path = os.path.join(inputs_dir, label)
        if label.startswith(".") or not os.path.isdir(path):
            continue
        for i, img in enumerate(os.listdir(path)):
            if i >= files_limit:
                break
            if img.startswith("."):
                continue
            img_path = os.path.join(path, img)
            input_files.append(img_path)
            input_labels.append(label)

    emotions = []
    for img in input_files:
        frame_bgr = cv2.imread(img)
        frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

        emotion, _ = fer.predict_emotions(frame, logits=True)
        emotions.append(emotion[0])

    assert len(emotions) == len(input_labels)
    preds = np.array(emotions)
    labels = np.array(input_labels)
    acc = (labels == preds).mean()
    assert acc > 0.55


@pytest.mark.parametrize("model_name", get_model_list())
@pytest.mark.parametrize("engine", ["torch", "onnx"])
def test_on_video(model_name, engine):
    """
    Simple test that checks emotions on video
    """
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    input_file = os.path.join(FILE_DIR, "data", "video_samples", "emotions", "Angry", "Angry.mp4")

    fer = EmotiEffLibRecognizer(engine=engine, model_name=model_name, device=device)

    cap = cv2.VideoCapture(input_file)
    all_scores = None
    while cap.isOpened():
        success, image = cap.read()
        if not success:
            break

        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        facial_images = recognize_faces(image_rgb, device)
        if len(facial_images) == 0:
            continue
        _, scores = fer.predict_emotions(facial_images, logits=True)
        if all_scores is not None:
            all_scores = np.concatenate((all_scores, scores))
        else:
            all_scores = scores

    cap.release()

    score = np.mean(all_scores, axis=0)
    emotion = np.argmax(score)

    assert fer.idx_to_emotion_class[emotion] == "Anger"


@pytest.mark.parametrize("model_name", get_model_list())
@pytest.mark.parametrize("engine", ["torch", "onnx"])
def test_engagement_on_video(model_name, engine):
    """
    Simple test that checks engagement on video
    """
    if model_name in ("enet_b2_8", "enet_b2_7"):
        pytest.xfail("These models are not supported")
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    input_file = os.path.join(
        FILE_DIR, "data", "video_samples", "engagement", "engaged", "1_video1.mp4"
    )

    fer = EmotiEffLibRecognizer(engine=engine, model_name=model_name, device=device)

    cap = cv2.VideoCapture(input_file)
    frames = []
    while cap.isOpened():
        success, image = cap.read()
        if not success:
            break

        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        facial_images = recognize_faces(image_rgb, device)
        if len(facial_images) == 0:
            continue
        frames += facial_images

    cap.release()
    _, scores = fer.predict_engagement(frames)

    score = np.mean(scores, axis=0)
    engagement = np.argmax(score)

    assert fer.idx_to_engagement_class[engagement] == "Engaged"


@pytest.mark.parametrize("model_name", get_model_list())
@pytest.mark.parametrize("engine", ["torch", "onnx"])
def test_distraction_on_video(model_name, engine):
    """
    Simple test that checks distraction on video
    """
    if model_name in ("enet_b2_8", "enet_b2_7"):
        pytest.xfail("These models are not supported")
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    input_file = os.path.join(
        FILE_DIR, "data", "video_samples", "engagement", "distracted", "0_video1.mp4"
    )

    fer = EmotiEffLibRecognizer(engine=engine, model_name=model_name, device=device)

    cap = cv2.VideoCapture(input_file)
    frames = []
    while cap.isOpened():
        success, image = cap.read()
        if not success:
            break

        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        facial_images = recognize_faces(image_rgb, device)
        if len(facial_images) == 0:
            continue
        frames += facial_images

    cap.release()
    _, scores = fer.predict_engagement(frames)

    score = np.mean(scores, axis=0)
    engagement = np.argmax(score)

    assert fer.idx_to_engagement_class[engagement] == "Distracted"


@pytest.mark.parametrize("model_name", get_model_list())
@pytest.mark.parametrize("engine", ["torch", "onnx"])
def test_engagement_and_emotion_on_video(model_name, engine):
    """
    Simple test that checks level of engagement and emotions on a video
    """
    if model_name in ("enet_b2_8", "enet_b2_7"):
        pytest.xfail("These models are not supported")
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    input_file = os.path.join(
        FILE_DIR, "data", "video_samples", "engagement", "engaged", "1_video1.mp4"
    )

    fer = EmotiEffLibRecognizer(engine=engine, model_name=model_name, device=device)

    cap = cv2.VideoCapture(input_file)
    frames = []
    while cap.isOpened():
        success, image = cap.read()
        if not success:
            break

        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        facial_images = recognize_faces(image_rgb, device)
        if len(facial_images) == 0:
            continue
        frames += facial_images

    cap.release()
    features = fer.extract_features(frames)
    _, emo_scores = fer.classify_emotions(features)
    _, eng_scores = fer.classify_engagement(features)

    emo_score = np.mean(emo_scores, axis=0)
    eng_score = np.mean(eng_scores, axis=0)
    emotion = np.argmax(emo_score)
    engagement = np.argmax(eng_score)

    assert fer.idx_to_engagement_class[engagement] == "Engaged"
    assert fer.idx_to_emotion_class[emotion] == "Sadness"
