"""FFHQ-512 face alignment, identical to ALADDIN's preprocessing.

Matches the face preprocessing of the ALADDIN code (https://github.com/isdevnull/aladdin),
which wraps facexlib 0.3.0 ``FaceRestoreHelper``:
RetinaFace-R50 detection, largest face only, 5-point similarity transform to the FFHQ 512
template (cv2.LMEDS), bilinear warpAffine with a constant gray border (135, 133, 132).

Needs ``pip install facexlib opencv-python``. Without ``model_root`` facexlib downloads
``detection_Resnet50_Final.pth`` from its GitHub release on first use.
"""

import cv2
import numpy as np

SIZE = 512
BORDER_VALUE = (135, 133, 132)


def read_bgr(path):
    """cv2.imread as ALADDIN does (applies EXIF orientation) -> uint8 BGR [H, W, 3]."""
    return cv2.imread(str(path), cv2.IMREAD_COLOR)


def warp(image, affine):
    """Crop with a known affine (original -> crop), exactly as facexlib's align_warp_face.

    Args:
        image: uint8 BGR [H, W, 3] original frame.
        affine: [2, 3] matrix.

    Returns:
        uint8 BGR [SIZE, SIZE, 3].
    """
    return cv2.warpAffine(image, np.asarray(affine, np.float64), (SIZE, SIZE),
                          borderMode=cv2.BORDER_CONSTANT, borderValue=BORDER_VALUE)


class FaceAligner:
    """RetinaFace-R50 + FFHQ-512 five-point alignment (state: the facexlib helper).

    Attributes:
        helper: facexlib FaceRestoreHelper, face_size 512, no parsing.
        min_eye_distance: float, detections with fewer pixels between the eyes are dropped.
    """

    def __init__(self, model_root=None, device="cuda", min_eye_distance=5.0):
        import torch
        from facexlib.utils.face_restoration_helper import FaceRestoreHelper

        self.min_eye_distance = min_eye_distance
        kwargs = {"model_rootpath": str(model_root)} if model_root else {}
        self.helper = FaceRestoreHelper(upscale_factor=1, face_size=SIZE, det_model="retinaface_resnet50",
                                        use_parse=False, device=torch.device(device), **kwargs)

    def align(self, image):
        """Largest face of a BGR frame -> (crop uint8 BGR [SIZE, SIZE, 3], affine [2, 3]) or None."""
        h = self.helper
        h.clean_all()
        h.read_image(image)
        h.get_face_landmarks_5(only_keep_largest=True, eye_dist_threshold=self.min_eye_distance)
        if not h.all_landmarks_5:
            return None
        h.align_warp_face(border_mode="constant")
        return h.cropped_faces[0], np.asarray(h.affine_matrices[0], np.float64)