"""Face blending module for MuseTalk"""

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


def get_image(image, face, face_box, mode="jaw", fp=None):
    """Blend generated face back to original image

    Args:
        image: Original full image (BGR numpy array)
        face: Generated face (BGR numpy array)
        face_box: Bounding box [x1, y1, x2, y2]
        mode: Blending mode ("jaw" or default)
        fp: Face parser (optional)
    """
    body = Image.fromarray(image[:, :, ::-1])
    face = Image.fromarray(face[:, :, ::-1])

    x1, y1, x2, y2 = face_box

    crop_box, s = _get_crop_box(face_box, expand=1.5)
    x_s, y_s, x_e, y_e = crop_box
    face_large = body.crop(crop_box)

    ori_shape = face_large.size

    if fp is not None:
        mask_image = fp(face_large, mode=mode)
        mask_small = mask_image.crop((x1 - x_s, y1 - y_s, x2 - x_s, y2 - y_s))
        mask_image = Image.new("L", ori_shape, 0)
        mask_image.paste(mask_small, (x1 - x_s, y1 - y_s, x2 - x_s, y2 - y_s))

        width, height = mask_image.size
        top_boundary = int(height * 0.5)
        modified_mask_image = Image.new("L", ori_shape, 0)
        modified_mask_image.paste(
            mask_image.crop((0, top_boundary, width, height)), (0, top_boundary)
        )

        blur_kernel_size = int(0.1 * ori_shape[0] // 2 * 2) + 1
        mask_array = cv2.GaussianBlur(
            np.array(modified_mask_image), (blur_kernel_size, blur_kernel_size), 0
        )
        mask_image = Image.fromarray(mask_array)
    else:
        mask_image = Image.new("L", ori_shape, 255)

    face_large.paste(face, (x1 - x_s, y1 - y_s, x2 - x_s, y2 - y_s))
    body.paste(face_large, crop_box[:2], mask_image)

    body = np.array(body)[:, :, ::-1]
    return body


def _get_crop_box(box, expand=1.5):
    x, y, x1, y1 = box
    x_c, y_c = (x + x1) // 2, (y + y1) // 2
    w, h = x1 - x, y1 - y
    s = int(max(w, h) // 2 * expand)
    crop_box = [x_c - s, y_c - s, x_c + s, y_c + s]
    return crop_box, s


if __name__ == "__main__":
    test_img = np.zeros((512, 512, 3), dtype=np.uint8)
    test_img[:, :, 1] = 100
    test_img[:, :, 2] = 200

    test_face = np.zeros((256, 256, 3), dtype=np.uint8)
    test_face[:, :, 0] = 255

    result = get_image(test_img, test_face, [128, 128, 384, 384], mode="jaw")
    cv2.imwrite("test_blending.png", result)
    print("Face blending test complete")