import numpy as np
import cv2
import math

# import blazeface


def euclideanDistance(a, b):
    """
    Calculates Euclidean distance.

    Parameters
    ----------
    a (and b): numpy array
        point in 2D space defined by [x, y] coordinates

    Return
    ------
    euclidean_distance: float
        calculated Euclidean distance
    """
    return math.sqrt(((b[0] - a[0]) ** 2) + ((b[1] - a[1]) ** 2))


class FaceAligner:
    """
    Aligns face image in reference to left and right eye coordinates, assuming roll is not extreme.

    Methods
    -------
    align(img, keypoints)
       aligns face image by rotating it to have left and right on same horizontal line

    testSelf(self)
        tests FaceAligner class
    """

    @staticmethod
    def align(img, keypoints):
        """
        Aligns face image by rotating it to have left and right eyes on same horizontal line.

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

        keypoints: list
            a list of left eye coordinates followed by right eye coordinates, each contains [x, y]
            coordinates normalized to the image width and height

        Return
        ------
        img_aligned: numpy array
            aligned (rotated) square cropped BGR image
        """
        # Build triangle based on 2 eyes and a 3rd point, 2 cases: left eye (keypoints[0]) is above
        # or below right eye (keypoints[1]), distance of y starts from top, and it's normalized to
        # the height of the image, so when left is lower:
        if keypoints[0][1] > keypoints[1][1]:
            point_3rd = (keypoints[1][0], keypoints[0][1])
            direction = -1
            # print("rotate to clock wise direction")
        else:
            point_3rd = (keypoints[0][0], keypoints[1][1])
            direction = 1
            # print("rotate to counter clock wise direction")

        # Calculate sides of the current triangle
        a = euclideanDistance(keypoints[0], point_3rd)
        b = euclideanDistance(keypoints[1], keypoints[0])
        c = euclideanDistance(keypoints[1], point_3rd)

        # Using cosine rule get wanted angle
        cos_a = (b * b + c * c - a * a) / (2 * b * c)
        angle = np.arccos(cos_a)
        angle = (angle * 180) / math.pi

        # In case the direction is clock wise need the completing angle
        if direction == -1:
            angle = 90 - angle

        # Rotate image center based on direction AND angle (scale stays same=1.0)!
        image_center = tuple(np.array(img.shape[1::-1]) / 2)
        rot_mat = cv2.getRotationMatrix2D(image_center, direction * angle, 1.0)
        # wtf - keep INTER_LINEAR for speed?
        img_aligned = cv2.warpAffine(
            img, rot_mat, img.shape[1::-1], flags=cv2.INTER_LINEAR
        )

        return img_aligned

    def testSelf(self, img, keypoints, precomputed_hash):
        """
        Tests FaceAligner class.

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

        keypoints: list
            a list of left eye coordinates followed by right eye coordinates, each contains [x, y]
            coordinates normalized to the image width and height

        precomputed_hash: int
            precomputed hash of the aligned image
        """
        # # Use BlazeFace to get keypoints
        # results = blazeface.BlazeFaceDetector(type="back").detectFaces(img)

        aligned_image = self.align(img, keypoints)

        # # For VISUAL debug:
        # # =================
        # dimensions = img.shape
        # img = cv2.line(img,
        #                (round(dimensions[1] * keypoints[0][0]),
        #                 round(dimensions[0] * keypoints[0][1])),
        #                (round(dimensions[1] * keypoints[1][0]),
        #                 round(dimensions[0] * keypoints[1][1])),
        #                color=(0, 0, 255),
        #                thickness=8)
        # cv2.imshow("pre align", img)
        # cv2.imshow("post align", aligned_image)
        # cv2.waitKey(0)
        # # =================

        # # Verify eyes are pretty much on same y-axis after alignment with BlazeFace
        # aligned_image_results = blazeface.BlazeFaceDetector(type="back").detectFaces(aligned_image)
        # assert abs(aligned_image_results.keypoints[0][0][1] - aligned_image_results.keypoints[0][1][1]) < 0.01

        # Need to precompute the hash and give as input to testSelf!
        aligned_image_hash = hash(tuple(aligned_image.flatten()))
        assert (
            precomputed_hash == aligned_image_hash
        ), f"aligned image hash: {aligned_image_hash}"

        return True
