from face_matcher.preprocessing.blazeface import BlazeFace


class FaceDetector:
    """
    Detects face(s) in image and returns score, keypoints, and bounding box.

    Methods
    -------
    detect(self, img)
        detects face(s) in image and returns score, keypoints, and bounding box
    """

    def __init__(self, model="back", postprocess="weighted_mean"):
        """
        Constructs face detector

        Parameters
        ----------
        model: string
            BlazeFace tflite model to be used, available models are "front", "back", and "short_range"

        postprocess: string
            method to be used for anchor postprocessing, either "non_max_suppression" or "weighted_mean"
        """
        self.model = model
        self.postprocess = postprocess
        self.face_detector = BlazeFace(model=self.model, postprocess=self.postprocess)

    def detect(self, img):
        """
        detects face(s) in image and returns score, keypoints, and bounding box.

        Parameters
        ----------
        img: numpy array
            square input BGR image

        Return
        ------
        scores: list
            BlazeFace confidence scores of detected faces

        keypoints: list
            eyes, ears, nose, and mouth location of detected faces

        boxes: list
            coordinates of bounding boxes of detected faces
        """
        results = self.face_detector.detect(img)
        scores = results.scores
        keypoints = results.keypoints
        boxes = results.boxes
        return scores, keypoints, boxes

    def testSelf(self, img):
        """
        Tests FaceDetector class.

        Parameters
        ----------
        img: numpy array
            square test BGR image
        """
        scores, keypoints, boxes = self.detect(img)

        assert abs(scores[0] - 0.9069122076034546) < 1e-6, f"score: {scores[0]}"

        assert abs(boxes[0][0] - 0.28554990098477506) < 1e-6, f"boxes x1: {boxes[0][0]}"
        assert abs(boxes[0][1] - 0.3907631600586229) < 1e-6, f"boxes y1: {boxes[0][1]}"

        assert abs(boxes[0][2] - 0.7501593515217089) < 1e-6, f"boxes x2: {boxes[0][2]}"
        assert abs(boxes[0][3] - 0.8553710561556538) < 1e-6, f"boxes y2: {boxes[0][3]}"

        assert (
            abs(keypoints[0][0][0] - 0.42408981804774953) < 1e-6
        ), f"keypoints left eye x: {keypoints[0][0][0]}"
        assert (
            abs(keypoints[0][0][1] - 0.5209796978768542) < 1e-6
        ), f"keypoints left eye y: {keypoints[0][0][1]}"

        assert (
            abs(keypoints[0][1][0] - 0.6132438914339988) < 1e-6
        ), f"keypoints right eye x: {keypoints[0][1][0]}"
        assert (
            abs(keypoints[0][1][1] - 0.5134623604214229) < 1e-6
        ), f"keypoints right eye y: {keypoints[0][1][1]}"

        assert (
            abs(keypoints[0][2][0] - 0.5273885515141743) < 1e-6
        ), f"keypoints nose x: {keypoints[0][2][0]}"
        assert (
            abs(keypoints[0][2][1] - 0.6296896756093392) < 1e-6
        ), f"keypoints nose y: {keypoints[0][2][1]}"

        assert (
            abs(keypoints[0][3][0] - 0.5293438327441915) < 1e-6
        ), f"keypoints mouth x: {keypoints[0][3][0]}"
        assert (
            abs(keypoints[0][3][1] - 0.7305445547360915) < 1e-6
        ), f"keypoints mouth y: {keypoints[0][3][1]}"

        assert (
            abs(keypoints[0][4][0] - 0.3240732091647567) < 1e-6
        ), f"keypoints left ear x: {keypoints[0][4][0]}"
        assert (
            abs(keypoints[0][4][1] - 0.5703979044571806) < 1e-6
        ), f"keypoints left ear y: {keypoints[0][4][1]}"

        assert (
            abs(keypoints[0][5][0] - 0.7065206648635853) < 1e-6
        ), f"keypoints right ear x: {keypoints[0][5][0]}"
        assert (
            abs(keypoints[0][5][1] - 0.555401002007089) < 1e-6
        ), f"keypoints right ear y: {keypoints[0][5][1]}"

        return True
