import math

import cv2
import numpy as np
import onnxruntime
from PIL import Image

session = onnxruntime.InferenceSession("models/slim-facelandmark.onnx")


def EuclideanDistance(source_representation, test_representation):
    euclidean_distance = source_representation - test_representation
    euclidean_distance = np.sum(np.multiply(euclidean_distance, euclidean_distance))
    euclidean_distance = np.sqrt(euclidean_distance)
    return euclidean_distance


def alignment_procedure(img, left_eye, right_eye):
    left_eye_x, left_eye_y = left_eye
    right_eye_x, right_eye_y = right_eye

    if left_eye_y > right_eye_y:
        point_3rd = (right_eye_x, left_eye_y)
        direction = -1  # rotate same direction to clock
    else:
        point_3rd = (left_eye_x, right_eye_y)
        direction = 1  # rotate inverse direction of clock

    a = EuclideanDistance(np.array(left_eye), np.array(point_3rd))
    b = EuclideanDistance(np.array(right_eye), np.array(point_3rd))
    c = EuclideanDistance(np.array(right_eye), np.array(left_eye))

    if (
        b != 0 and c != 0
    ):  # this multiplication causes division by zero in cos_a calculation

        cos_a = (b * b + c * c - a * a) / (2 * b * c)
        angle = np.arccos(cos_a)  # angle in radian
        angle = (angle * 180) / math.pi  # radian to degree

        # -----------------------
        # rotate base image

        if direction == -1:
            angle = 90 - angle

        img = Image.fromarray(img)
        img = np.array(img.rotate(direction * angle))

    # -----------------------

    return img  # return img anyway


def align_face(image):
    inputs = cv2.resize(image, (112, 112))
    inputs = cv2.cvtColor(inputs, cv2.COLOR_BGR2RGB)
    inputs = inputs.transpose(2, 0, 1).astype(np.float32)
    inputs = inputs / 255.0
    inputs = np.expand_dims(inputs, axis=0)
    landmarks = session.run(None, {"input": inputs})[0]
    pre_landmark = landmarks[0]
    pre_landmark = pre_landmark.reshape(-1, 2)
    left_eyex, left_eyey = pre_landmark[96]
    right_eyex, right_eyey = pre_landmark[97]

    img = alignment_procedure(image, (left_eyex, left_eyey), (right_eyex, right_eyey))
    return img


if __name__ == "__main__":
    img = cv2.imread("obamaface.jpg")
    img = align_face(img)
    cv2.imshow("img", img)
    cv2.waitKey(0)