from mtcnn import MTCNN
from facenet_pytorch import MTCNN
import torch
import math
import cv2
import numpy as np

K_MULTIPLIER = 1.2

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
detector = MTCNN(keep_all=True, device=device)


class FaceDetection:
    def __init__(self,):
        self.detector = detector

    def calculate_distance(self, p1, p2) -> float:
        return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)

    def calculate_dis_to_cp(self, cx, cy, face_cx, face_cy) -> float:
        return math.sqrt((face_cx - cx) ** 2 + (face_cy - cy) ** 2)

    def detect_face(self, frame, cx, cy) -> bool:
        boxes, probs, landmarks = self.detector.detect(frame, landmarks=True)

        for box, landmark in zip(boxes, landmarks):
            # Draw bounding box
            x1, y1, x2, y2 = map(int, box)

            face_cx = int(x1 + (x2 - x1) / 2)
            face_cy = int(y1 + (y2 - y1) / 2)

            if len(landmark) >= 5:
                nose = landmark[2]
                left_eye = landmark[0]
                right_eye = landmark[1]

                # Calculate distances
                distance_left = self.calculate_distance(nose, left_eye)
                distance_right = self.calculate_distance(nose, right_eye)

                # Check if distances exceed threshold
                if not (distance_left > K_MULTIPLIER * distance_right or distance_right > K_MULTIPLIER * distance_left or
                        self.calculate_dis_to_cp(cx, cy, face_cx, face_cy) > 30):
                    return True

        return False