import numpy as np
import onnxruntime as rt
import mediapipe as mp
import cv2
import os
import time
from skimage.transform import SimilarityTransform


# ---------------------------------------------------------------------------------------------------------------------
# INITIALIZATIONS

# Target landmark coordinates for alignment (used in training)
LANDMARKS_TARGET = np.array(
    [
        [38.2946, 51.6963],
        [73.5318, 51.5014],
        [56.0252, 71.7366],
        [41.5493, 92.3655],
        [70.7299, 92.2041],
    ],
    dtype=np.float32,
)

# Initialize Face Detector (For Example Mediapipe)
FACE_DETECTOR = mp.solutions.face_mesh.FaceMesh(
    refine_landmarks=True, min_detection_confidence=0.5, min_tracking_confidence=0.5, max_num_faces=1
)

# Initialize the Face Recognition Model (FaceTransformerOctupletLoss)
FACE_RECOGNIZER = rt.InferenceSession("FaceTransformerOctupletLoss.onnx", providers=rt.get_available_providers())


# ---------------------------------------------------------------------------------------------------------------------
# FACE CAPTURE

# Capture a frame with your Webcam and store it on disk
if not os.path.exists("img.jpg"):
    cap = cv2.VideoCapture(1) # open webcam
    time.sleep(2) # wait for camera to warm up
    
    if not cap.isOpened():
        raise IOError("Cannot open webcam")
    
    ret, img = cap.read() # capture a frame
    if ret:
        cv2.imwrite("img.jpg", img) # save the frame
else:
    img = cv2.imread("img.jpg") # read the frame from disk


# ---------------------------------------------------------------------------------------------------------------------
# FACE DETECTION

# Process the image with the face detector
result = FACE_DETECTOR.process(img)

if result.multi_face_landmarks:
    # Select 5 Landmarks (Eye Centers, Nose Tip, Left Mouth Corner, Right Mouth Corner)
    five_landmarks = np.asarray(result.multi_face_landmarks[0].landmark)[[470, 475, 1, 57, 287]]

    # Extract the x and y coordinates of the landmarks of interest
    landmarks = np.asarray(
        [[landmark.x * img.shape[1], landmark.y * img.shape[0]] for landmark in five_landmarks]
    )

    # Extract the x and y coordinates of all landmarks
    all_x_coords = [landmark.x * img.shape[1] for landmark in result.multi_face_landmarks[0].landmark]
    all_y_coords = [landmark.y * img.shape[0] for landmark in result.multi_face_landmarks[0].landmark]

    # Compute the bounding box of the face
    x_min, x_max = int(min(all_x_coords)), int(max(all_x_coords))
    y_min, y_max = int(min(all_y_coords)), int(max(all_y_coords))
    bbox = [[x_min, y_min], [x_max, y_max]]

else:
    print("No faces detected")
    exit()


# ---------------------------------------------------------------------------------------------------------------------
# FACE ALIGNMENT

# Align Image with the 5 Landmarks
tform = SimilarityTransform()
tform.estimate(landmarks, LANDMARKS_TARGET)
tmatrix = tform.params[0:2, :]
img_aligned = cv2.warpAffine(img, tmatrix, (112, 112), borderValue=0.0)

# safe to disk
cv2.imwrite("img2_aligned.jpg", img_aligned)


# ---------------------------------------------------------------------------------------------------------------------
# FACE RECOGNITION

# Inference face embeddings with onnxruntime
input_image = (np.asarray([img_aligned]).astype(np.float32)).clip(0.0, 255.0).transpose(0, 3, 1, 2)
embedding = FACE_RECOGNIZER.run(None, {"input_image": input_image})[0][0]

print("Embedding:", embedding)

# If you have embeddings for several facial images - you can then compute the cosine distance between them and distinguish
# between different or same people based on a threshold. For example, if the cosine distance is less than 0.5, then the
# two images are of the same person, otherwise they are of different people. The lower the cosine distance, the more similar
# the two images are. The cosine distance is a value between 0 and 2, where 0 means the two images are identical and 2 means 
# the two images are completely different. 

# ---------------------------------------------------------------------------------------------------------------------
# VISUALIZATION

# Draw Boundingbox on a copy of image
img_draw = img.copy()
cv2.rectangle(img_draw, (bbox[0][0], bbox[0][1]), (bbox[1][0], bbox[1][1]), (255, 0, 0), 2)

# Show the detected face on the image
cv2.imshow("img", img_draw)
cv2.waitKey(0)

# Show the aligned image
cv2.imshow("img", img_aligned)
cv2.waitKey(0)