from codeformer.facelib.utils.face_restoration_helper import FaceRestoreHelper

import numpy as np
from codeformer.basicsr.utils.misc import get_device

class FaceRestoreHelperOptimized(FaceRestoreHelper):
    def __init__(
            self,
            upscale_factor,
            face_size=512,
            crop_ratio=(1, 1),
            det_model="retinaface_resnet50",
            save_ext="png",
            template_3points=False,
            pad_blur=False,
            use_parse=False,
            device=None,
    ):
        self.template_3points = template_3points  # improve robustness
        self.upscale_factor = int(upscale_factor)
        # the cropped face ratio based on the square face
        self.crop_ratio = crop_ratio  # (h, w)
        assert self.crop_ratio[0] >= 1 and self.crop_ratio[1] >= 1, "crop ration only supports >=1"
        self.face_size = (int(face_size * self.crop_ratio[1]), int(face_size * self.crop_ratio[0]))
        self.det_model = det_model

        if self.det_model == "dlib":
            # standard 5 landmarks for FFHQ faces with 1024 x 1024
            self.face_template = np.array(
                [
                    [686.77227723, 488.62376238],
                    [586.77227723, 493.59405941],
                    [337.91089109, 488.38613861],
                    [437.95049505, 493.51485149],
                    [513.58415842, 678.5049505],
                ]
            )
            self.face_template = self.face_template / (1024 // face_size)
        elif self.template_3points:
            self.face_template = np.array([[192, 240], [319, 240], [257, 371]])
        else:
            # standard 5 landmarks for FFHQ faces with 512 x 512
            # facexlib
            self.face_template = np.array(
                [
                    [192.98138, 239.94708],
                    [318.90277, 240.1936],
                    [256.63416, 314.01935],
                    [201.26117, 371.41043],
                    [313.08905, 371.15118],
                ]
            )

            # dlib: left_eye: 36:41  right_eye: 42:47  nose: 30,32,33,34  left mouth corner: 48  right mouth corner: 54
            # self.face_template = np.array([[193.65928, 242.98541], [318.32558, 243.06108], [255.67984, 328.82894],
            #                                 [198.22603, 372.82502], [313.91018, 372.75659]])

        self.face_template = self.face_template * (face_size / 512.0)
        if self.crop_ratio[0] > 1:
            self.face_template[:, 1] += face_size * (self.crop_ratio[0] - 1) / 2
        if self.crop_ratio[1] > 1:
            self.face_template[:, 0] += face_size * (self.crop_ratio[1] - 1) / 2
        self.save_ext = save_ext
        self.pad_blur = pad_blur
        if self.pad_blur is True:
            self.template_3points = False

        self.all_landmarks_5 = []
        self.det_faces = []
        self.affine_matrices = []
        self.inverse_affine_matrices = []
        self.cropped_faces = []
        self.restored_faces = []
        self.pad_input_imgs = []

        if device is None:
            # self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            self.device = get_device()
        else:
            self.device = device

        # init face detection model
        # if self.det_model == "dlib":
        #     self.face_detector, self.shape_predictor_5 = self.init_dlib(
        #         dlib_model_url["face_detector"], dlib_model_url["shape_predictor_5"]
        #     )
        # else:
        #     self.face_detector = init_detection_model(det_model, half=False, device=self.device)

        # init face parsing model
        self.use_parse = use_parse
        #self.face_parse = init_parsing_model(model_name="parsenet", device=self.device)

        # MUST set face_detector and face_parse!!!