from typing import Union
from pitop.processing.core.load_models import load_face_landmark_predictor
from pitop.core import ImageFunctions
from imutils import (
    face_utils,
    resize,
)
from imutils.video import FPS
from os import getenv
import atexit
from pitop.processing.core.vision_functions import (
    import_dlib,
    import_opencv,
    tuple_for_color_by_name,
)
from .core.face import Face


cv2 = import_opencv()
dlib = import_dlib()


class FaceDetector:
    _FACE_DETECTOR_PYRAMID_LAYERS = 1  # set higher to detect smaller faces. Cost: large reduction in detection FPS.

    def __init__(self,
                 image_processing_width: Union[int, None] = 320,
                 format: str = "OpenCV",
                 enable_tracking: bool = True,
                 dlib_landmark_predictor_filename: str = "shape_predictor_68_face_landmarks.dat"
                 ):
        """
        :param Union[int, None] image_processing_width: image width to scale to for image processing, set to None for
        no scaling.
        :param str format: desired output image format.
        :param bool enable_tracking: enable dlib's correlaction tracker to track the detected face between frames.
        :param str dlib_landmark_predictor_filename: Filename for facial features predictor. Use 5 landmark version for
        slightly better performance whilst retaining ability to calculate face angle. May be incompatible with further
         processing e.g. emotion detection requires the 68-landmark version.
        """
        self._image_processing_width = image_processing_width
        self._format = format
        self._face_rectangle_detector = dlib.get_frontal_face_detector()
        self._predictor = load_face_landmark_predictor(model_filename=dlib_landmark_predictor_filename)
        self._clahe_filter = cv2.createCLAHE(clipLimit=5)
        self._frame_scaler = None
        self.face = Face()
        self._enable_tracking = enable_tracking
        self._face_tracker = None

        # Enable FPS if environment variable is set
        self._print_fps = getenv("PT_ENABLE_FPS", "0") == "1"
        if self._print_fps:
            self._fps = FPS().start()
            atexit.register(self.__print_fps)

    def __call__(self, frame):
        """Detect a face in an image frame.

        :param frame: Image frame in OpenCV or PIL format.
        :return: Face object containing data about the detected face.
        """
        frame = ImageFunctions.convert(frame, format='OpenCV')

        if self._frame_scaler is None:
            _, width = frame.shape[0:2]
            self._frame_scaler = width / self._image_processing_width if self._image_processing_width is not None else 1

        frame_to_process = self.__get_frame_to_process(frame=frame)

        if self._face_tracker is None:
            # if the face tracker has not been started, we first need to use face detection to find a face
            face_rectangle, face_center, face_features = self.__detect_largest_face(frame=frame_to_process)
            if face_center is not None and self._enable_tracking:
                # Face found, enable dlib correlation tracker for subsequent calls
                self.__start_tracker(frame=frame_to_process, rectangle=face_rectangle)
        else:
            # We are in face tracking mode, use dlib correlation tracker to track the face
            face_rectangle, face_center, face_features = self.__track_face(frame=frame_to_process)
            if face_center is None:
                self.__stop_tracker()
                # attempt to detect face since tracker has failed
                face_rectangle, face_center, face_features = self.__detect_largest_face(frame=frame_to_process)

        self.face = self.__prepare_face_data(frame=frame,
                                             face=self.face,
                                             rectangle=face_rectangle,
                                             center=face_center,
                                             features=face_features)
        if self._print_fps:
            self._fps.update()

        return self.face

    def __get_frame_to_process(self, frame):
        """Resize frame, convert to grayscale and use contrast limited adaptive
        histogram equalization (CLAHE) to improve the contrast of the image so
        that face features are more pronounced.

        :param frame: original frame from the camera in OpenCV format
        :return: OpenCV image to send to face processing algorithms
        """
        return self._clahe_filter.apply(
            cv2.cvtColor(
                resize(frame, width=self._image_processing_width),
                cv2.COLOR_BGR2GRAY
            )
        )

    def __detect_largest_face(self, frame):
        """Use dlib face rectangle detector on the frame and process the
        detected rectangles to get the face data.

        :param frame: Image frame to process for faces
        :return: Detected face data
        """
        face_rectangle, face_center, face_features = self.__process_detected_rectangles(
            frame,
            self._face_rectangle_detector(frame, self._FACE_DETECTOR_PYRAMID_LAYERS)
        )

        return face_rectangle, face_center, face_features

    def __process_detected_rectangles(self, frame, rectangles_dlib):
        """Find largest face rectangle and process it to get the required face
        data.

        :param frame: OpenCV frame to use for processing
        :param rectangles_dlib: Rectangles found using dlib's face detector in the dlib Rectangle format.
        :return: Detected face data
        """
        if len(rectangles_dlib) == 0:
            return None, None, None

        largest_rectangle_dlib = None
        area = 0
        face_rectangle = None
        face_center = None
        for rectangle_dlib in rectangles_dlib:
            x, y, w, h = face_utils.rect_to_bb(rectangle_dlib)
            current_area = w * h
            if current_area > area:
                area = current_area
                largest_rectangle_dlib = rectangle_dlib
                face_rectangle = (x, y, w, h)
                face_center = (int(x + w / 2), int(y + h / 2))

        face_features = self.__get_dlib_face_features(frame, largest_rectangle_dlib)

        return face_rectangle, face_center, face_features

    def __get_dlib_face_features(self, frame, dlib_rectangle):
        """Find's dlib face features using the predictor and then convert them
        to numpy format.

        :param frame: frame used to detect/track face rectangle.
        :param dlib_rectangle: Face rectangle in dlib Rectangle format.
        :return: 68x2 numpy array of x, y coordinates for facial features.
        """
        face_features_dlib = self._predictor(frame, dlib_rectangle)
        return face_utils.shape_to_np(face_features_dlib)

    def __start_tracker(self, frame, rectangle):
        self._face_tracker = dlib.correlation_tracker()
        self._face_tracker.start_track(frame,
                                       dlib.rectangle(
                                           rectangle[0],
                                           rectangle[1],
                                           rectangle[0] + rectangle[2],
                                           rectangle[1] + rectangle[3]
                                       )
                                       )

    def __stop_tracker(self):
        self._face_tracker = None

    def __track_face(self, frame):
        """Use dlib correlation tracker to track the face from the last frame
        in the current frame.

        :param frame: OpenCV frame to use for processing
        :return: Face data
        """
        def convert_dlib_rect_to_int_type(rectangle):
            x_start = int(round(rectangle.left()))
            y_start = int(round(rectangle.top()))
            x_end = int(round(rectangle.right()))
            y_end = int(round(rectangle.bottom()))
            return dlib.rectangle(x_start, y_start, x_end, y_end)

        peak_to_side_lobe_ratio = self._face_tracker.update(frame)

        if peak_to_side_lobe_ratio < 7.0:
            # Object occluded or lost when PSR found to be below a certain threshold (7.0 according to Bolme et al)
            # Treat as failed attempt to track face and return
            return None, None, None

        rectangle_dlib = convert_dlib_rect_to_int_type(self._face_tracker.get_position())
        face_features = self.__get_dlib_face_features(frame, rectangle_dlib)

        face_rectangle = face_utils.rect_to_bb(rectangle_dlib)
        face_center = (face_rectangle[0] + int(round(face_rectangle[2] / 2)),
                       face_rectangle[1] + int(round(face_rectangle[3] / 2)))

        return face_rectangle, face_center, face_features

    def __prepare_face_data(self, frame, face, rectangle, center, features):
        face.original_detection_frame = frame

        if center is None:
            face.clear()
            face.robot_view = ImageFunctions.convert(frame, format=self._format)
            return face

        # resize back to original frame resolution
        face.rectangle = tuple((int(item * self._frame_scaler) for item in rectangle))
        face.center_default = tuple((int(item * self._frame_scaler) for item in center))
        face.features = (features * self._frame_scaler).astype("int")

        face.robot_view = ImageFunctions.convert(
            self.__draw_on_frame(frame=frame.copy(), face=face),
            format=self._format
        )
        return face

    @staticmethod
    def __draw_on_frame(frame, face):
        x, y, w, h = face.rectangle
        cv2.rectangle(frame, (x, y), (x + w, y + h), tuple_for_color_by_name("dodgerblue", bgr=True), 2)

        cv2.drawMarker(frame, face.center_default, tuple_for_color_by_name("orangered", bgr=True),
                       markerType=cv2.MARKER_CROSS, markerSize=10, thickness=3, line_type=cv2.FILLED)

        return frame

    def __print_fps(self):
        self._fps.stop()
        print(f"[INFO] Elapsed time: {self._fps.elapsed():.2f}")
        print(f"[INFO] Approx. FPS: {self._fps.fps():.2f}")
