"""Face Parsing module for MuseTalk - Simplified implementation"""

import torch
import cv2
import numpy as np
from PIL import Image
import torchvision.transforms as transforms


class FaceParsing:
    def __init__(self, left_cheek_width=80, right_cheek_width=80):
        self.left_cheek_width = left_cheek_width
        self.right_cheek_width = right_cheek_width
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def __call__(self, image, size=(512, 512), mode="jaw"):
        if isinstance(image, str):
            image = Image.open(image)

        width, height = image.size

        with torch.no_grad():
            image = image.resize(size, Image.BILINEAR)

            face_mask = self._get_face_mask(image, mode=mode)

            return face_mask

    def _get_face_mask(self, image, mode="jaw"):
        width, height = image.size

        if mode == "jaw":
            mask = np.zeros((height, width), dtype=np.uint8)
            mask[:, int(height * 0.5) :] = 255
            kernel = np.ones((21, 21), dtype=np.uint8)
            kernel = cv2.erode(kernel, np.ones((3, 3), dtype=np.uint8), iterations=2)
            cheek_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (35, 3))
            mask = cv2.erode(mask, cheek_kernel, iterations=2)

            cheek_mask = np.zeros((height, width), dtype=np.uint8)
            center = width // 2
            cv2.rectangle(
                cheek_mask, (0, 0), (center - self.left_cheek_width, height), 255, -1
            )
            cv2.rectangle(
                cheek_mask,
                (center + self.right_cheek_width, 0),
                (width, height),
                255,
                -1,
            )
            mask = cv2.bitwise_and(mask, cheek_mask)

            return Image.fromarray(mask)
        else:
            mask = np.zeros((height, width), dtype=np.uint8)
            mask[:, int(height * 0.5) :] = 255
            return Image.fromarray(mask)


if __name__ == "__main__":
    fp = FaceParsing()
    test_img = Image.new("RGB", (512, 512), (100, 150, 200))
    mask = fp(test_img, mode="jaw")
    mask.save("test_mask.png")
    print("Face parsing test complete")