import os
import cv2
import dlib
import mediapipe as mp

def read_image(image_path):
    """
    Read an image from the given path and convert it to RGB.
    """
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found at path: {image_path}")
    
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

class LandmarksDetector:
    def __init__(self, predictor_model_path=f'{os.path.dirname(os.path.abspath(__file__))}/utils/shape_predictor_68_face_landmarks.dat'):
        """
        :param predictor_model_path: path to shape_predictor_68_face_landmarks.dat file
        """
        self.detector = dlib.get_frontal_face_detector() # cnn_face_detection_model_v1 also can be used
        self.shape_predictor = dlib.shape_predictor(predictor_model_path)

    def get_landmarks(self, image):
        img = dlib.load_rgb_image(image)
        dets = self.detector(img, 1)
        dets = [dets[0]]

        for detection in dets:
            try:
                face_landmarks = [(item.x, item.y) for item in self.shape_predictor(img, detection).parts()]
                yield face_landmarks
            except:
                print("Exception in get_landmarks()!")


class DlibLandmarkDetector:
    def __init__(self, predictor_model_path=f'{os.path.dirname(os.path.abspath(__file__))}/utils/shape_predictor_68_face_landmarks.dat'):
        """
        :param predictor_model_path: path to shape_predictor_68_face_landmarks.dat file
        """
        self.detector = dlib.get_frontal_face_detector() # cnn_face_detection_model_v1 also can be used
        self.predictor = dlib.shape_predictor(predictor_model_path)

    def get_landmarks(self, image_path):
        # image = dlib.load_rgb_image(image_path)
        image = read_image(image_path)
        height, width, _ = image.shape

        # Detect the faces in the image
        dets = self.detector(image, 1)  # 1 indicates to upsample the image 1 time. Higher values may give better results
        
        # Raise an exception if no face is detected
        if len(dets) == 0:
            raise Exception("No face detected in the image at path: ", image_path)
        
        # Get the landmarks of the first face detected
        face_landmarks = [(item.x, item.y) for item in self.predictor(image, dets[0]).parts()]

        # Add corner and edge midpoints as landmarks to include the background
        corner_landmarks = [(1, 1),  (1, height - 1), (width - 1, 1), (width - 1, height - 1)]
        edge_landmarks = [(1, (height - 1)//2), ((width - 1)//2, 1), ((width - 1)//2, height - 1), (width - 1, (height - 1)//2)]

        # Concatenate the landmarks
        landmarks = face_landmarks + corner_landmarks + edge_landmarks

        return landmarks, image

    def show_landmarked_image(self, image_path, landmarks):
        image = read_image(image_path)

        for landmark in landmarks:
            x, y = landmark
            cv2.circle(image, (x, y), 1, (255, 255, 0), -1) # image, (x, y), radius, color, thickness (-1 to fill)

        cv2.imshow('image', image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()


class  MediaPipeLandmarkDetector:
    def __init__(self):
 
        self.face_mesh = mp.solutions.face_mesh.FaceMesh(
            static_image_mode=True, 
            max_num_faces=1, 
            min_detection_confidence=0.5)


    def get_landmarks(self, image_path):

        image = read_image(image_path)
        height, width, _ = image.shape
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # Process the image
        results = self.face_mesh.process(image_rgb)

        # Raise an exception if no face is detected
        if results.multi_face_landmarks is None:
            raise Exception("No face detected in the image at path: ", image_path)
        
        # Extract the face landmarks
        face_landmarks = results.multi_face_landmarks[0]
        face_landmarks_normalized = [[landmark.x , landmark.y] for landmark in face_landmarks.landmark]

        # Add corner and edge midpoints as landmarks to include the background
        corner_landmarks = [(0, 0), (0, 1), (1, 0), (1, 1)]
        edge_landmarks = [(0, 0.5), (0.5, 0), (0.5, 1), (1, 0.5)]

        # Concatenate the corner and edge landmarks
        landmarks = corner_landmarks + edge_landmarks + face_landmarks_normalized
        
        # Multiply the landmarks with the image dimensions
        landmarks = [(int(x * width) - 1, int(y * height) - 1) for x, y in landmarks]
        landmarks = [(max(1, x), max(1, y)) for x, y in landmarks]

        return landmarks, image

    def show_landmarked_image(self, image_path, landmarks):
        image = cv2.imread(image_path)

        for landmark in landmarks:
            x, y = landmark
            cv2.circle(image, (x, y), 1, (255, 255, 0), -1)

        cv2.imshow('image', image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()