import cv2
import numpy as np
from insightface.model_zoo.inswapper import INSwapper
from insightface.utils import face_align
from modules import processing, shared
from modules.shared import opts
from modules.upscaler import UpscalerData

from scripts.faceswaplab_postprocessing import upscaling
from scripts.faceswaplab_postprocessing.postprocessing_options import (
    PostProcessingOptions,
)
from scripts.faceswaplab_swapping.facemask import generate_face_mask
from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2


def get_upscaler() -> UpscalerData:
    for upscaler in shared.sd_upscalers:
        if upscaler.name == opts.data.get(
            "faceswaplab_upscaled_swapper_upscaler", "LDSR"
        ):
            return upscaler
    return None


def merge_images_with_mask(image1, image2, mask):
    if image1.shape != image2.shape or image1.shape[:2] != mask.shape:
        raise ValueError("Img should have the same shape")
    mask = mask.astype(np.uint8)
    masked_region = cv2.bitwise_and(image2, image2, mask=mask)
    inverse_mask = cv2.bitwise_not(mask)
    empty_region = cv2.bitwise_and(image1, image1, mask=inverse_mask)
    merged_image = cv2.add(empty_region, masked_region)
    return merged_image


def erode_mask(mask, kernel_size=3, iterations=1):
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    eroded_mask = cv2.erode(mask, kernel, iterations=iterations)
    return eroded_mask


def apply_gaussian_blur(mask, kernel_size=(5, 5), sigma_x=0):
    blurred_mask = cv2.GaussianBlur(mask, kernel_size, sigma_x)
    return blurred_mask


def dilate_mask(mask, kernel_size=5, iterations=1):
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    dilated_mask = cv2.dilate(mask, kernel, iterations=iterations)
    return dilated_mask


def get_face_mask(aimg, bgr_fake):
    mask1 = generate_face_mask(aimg, device=shared.device)
    mask2 = generate_face_mask(bgr_fake, device=shared.device)
    mask = dilate_mask(cv2.bitwise_or(mask1, mask2))
    return mask


class UpscaledINSwapper:
    def __init__(self, inswapper: INSwapper):
        self.__dict__.update(inswapper.__dict__)

    def forward(self, img, latent):
        img = (img - self.input_mean) / self.input_std
        pred = self.session.run(
            self.output_names, {self.input_names[0]: img, self.input_names[1]: latent}
        )[0]
        return pred

    def super_resolution(self, img, k=2):
        pil_img = cv2_to_pil(img)
        options = PostProcessingOptions(
            upscaler_name=opts.data.get(
                "faceswaplab_upscaled_swapper_upscaler", "LDSR"
            ),
            upscale_visibility=1,
            scale=k,
            face_restorer_name=opts.data.get(
                "faceswaplab_upscaled_swapper_face_restorer", ""
            ),
            codeformer_weight=opts.data.get(
                "faceswaplab_upscaled_swapper_face_restorer_weight", 1
            ),
            restorer_visibility=opts.data.get(
                "faceswaplab_upscaled_swapper_face_restorer_visibility", 1
            ),
        )
        upscaled = upscaling.upscale_img(pil_img, options)
        upscaled = upscaling.restore_face(upscaled, options)
        return pil_to_cv2(upscaled)

    def get(self, img, target_face, source_face, paste_back=True, upscale=True):
        aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0])
        blob = cv2.dnn.blobFromImage(
            aimg,
            1.0 / self.input_std,
            self.input_size,
            (self.input_mean, self.input_mean, self.input_mean),
            swapRB=True,
        )
        latent = source_face.normed_embedding.reshape((1, -1))
        latent = np.dot(latent, self.emap)
        latent /= np.linalg.norm(latent)
        pred = self.session.run(
            self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent}
        )[0]
        # print(latent.shape, latent.dtype, pred.shape)
        img_fake = pred.transpose((0, 2, 3, 1))[0]
        bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:, :, ::-1]

        try:
            if not paste_back:
                return bgr_fake, M
            else:
                target_img = img

                def compute_diff(bgr_fake, aimg):
                    fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
                    fake_diff = np.abs(fake_diff).mean(axis=2)
                    fake_diff[:2, :] = 0
                    fake_diff[-2:, :] = 0
                    fake_diff[:, :2] = 0
                    fake_diff[:, -2:] = 0
                    return fake_diff

                if upscale:
                    print("*" * 80)
                    print(
                        f"Upscaled inswapper using {opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR')}"
                    )
                    print("*" * 80)

                    k = 4
                    aimg, M = face_align.norm_crop2(
                        img, target_face.kps, self.input_size[0] * k
                    )

                    # upscale and restore face :
                    bgr_fake = self.super_resolution(bgr_fake, k)

                    if opts.data.get("faceswaplab_upscaled_improved_mask", True):
                        mask = get_face_mask(aimg, bgr_fake)
                        bgr_fake = merge_images_with_mask(aimg, bgr_fake, mask)

                    # compute fake_diff before sharpen and color correction (better result)
                    fake_diff = compute_diff(bgr_fake, aimg)

                    if opts.data.get("faceswaplab_upscaled_swapper_sharpen", True):
                        print("sharpen")
                        # Add sharpness
                        blurred = cv2.GaussianBlur(bgr_fake, (0, 0), 3)
                        bgr_fake = cv2.addWeighted(bgr_fake, 1.5, blurred, -0.5, 0)

                    # Apply color corrections
                    if opts.data.get("faceswaplab_upscaled_swapper_fixcolor", True):
                        print("color correction")
                        correction = processing.setup_color_correction(cv2_to_pil(aimg))
                        bgr_fake_pil = processing.apply_color_correction(
                            correction, cv2_to_pil(bgr_fake)
                        )
                        bgr_fake = pil_to_cv2(bgr_fake_pil)

                else:
                    fake_diff = compute_diff(bgr_fake, aimg)

                IM = cv2.invertAffineTransform(M)

                img_white = np.full(
                    (aimg.shape[0], aimg.shape[1]), 255, dtype=np.float32
                )
                bgr_fake = cv2.warpAffine(
                    bgr_fake,
                    IM,
                    (target_img.shape[1], target_img.shape[0]),
                    borderValue=0.0,
                )
                img_white = cv2.warpAffine(
                    img_white,
                    IM,
                    (target_img.shape[1], target_img.shape[0]),
                    borderValue=0.0,
                )
                fake_diff = cv2.warpAffine(
                    fake_diff,
                    IM,
                    (target_img.shape[1], target_img.shape[0]),
                    borderValue=0.0,
                )
                img_white[img_white > 20] = 255
                fthresh = opts.data.get("faceswaplab_upscaled_swapper_fthresh", 10)
                print("fthresh", fthresh)
                fake_diff[fake_diff < fthresh] = 0
                fake_diff[fake_diff >= fthresh] = 255
                img_mask = img_white
                mask_h_inds, mask_w_inds = np.where(img_mask == 255)
                mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
                mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
                mask_size = int(np.sqrt(mask_h * mask_w))
                erosion_factor = opts.data.get(
                    "faceswaplab_upscaled_swapper_erosion", 1
                )
                k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor))

                kernel = np.ones((k, k), np.uint8)
                img_mask = cv2.erode(img_mask, kernel, iterations=1)
                kernel = np.ones((2, 2), np.uint8)
                fake_diff = cv2.dilate(fake_diff, kernel, iterations=1)
                k = max(int(mask_size // 20 * erosion_factor), int(5 * erosion_factor))

                kernel_size = (k, k)
                blur_size = tuple(2 * i + 1 for i in kernel_size)
                img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
                k = int(5 * erosion_factor)
                kernel_size = (k, k)
                blur_size = tuple(2 * i + 1 for i in kernel_size)
                fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0)
                img_mask /= 255
                fake_diff /= 255

                img_mask = np.reshape(
                    img_mask, [img_mask.shape[0], img_mask.shape[1], 1]
                )
                fake_merged = img_mask * bgr_fake + (1 - img_mask) * target_img.astype(
                    np.float32
                )
                fake_merged = fake_merged.astype(np.uint8)
                return fake_merged
        except Exception as e:
            import traceback

            traceback.print_exc()
            raise e