from confs import *
from pathlib import Path
import numpy as np
import cv2
from PIL import Image
import torch
import torch.utils.data as data
import torchvision.transforms as T
from einops import rearrange
import albumentations

from util_face import *
from util_4dataset import *
from util_cv2 import cv2_resize_auto_interpolation
from Mediapipe_Result_Cache import Mediapipe_Result_Cache


def resize_A(img, dataset_name, size=(512, 512), interpolation=None):
    is_pil = isinstance(img, Image.Image)
    if is_pil:
        img = np.array(img)
    if img.shape[:2] != (512, 512):
        img = cv2_resize_auto_interpolation(img, size, interpolation=interpolation)
    if is_pil:
        img = Image.fromarray(img)
    return img


def un_norm_clip(x1):
    x = x1 * 1.0
    reduce = False
    if len(x.shape) == 3:
        x = x.unsqueeze(0)
        reduce = True
    x[:, 0, :, :] = x[:, 0, :, :] * 0.26862954 + 0.48145466
    x[:, 1, :, :] = x[:, 1, :, :] * 0.26130258 + 0.4578275
    x[:, 2, :, :] = x[:, 2, :, :] * 0.27577711 + 0.40821073
    if reduce:
        x = x.squeeze(0)
    return x


def un_norm(x):
    return (x + 1.0) / 2.0


def _dilate(_mask, kernel_size, iterations):
    _mask = _mask.astype(np.uint8)
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    _mask = cv2.dilate(_mask, kernel, iterations=iterations)
    _mask = _mask.astype(bool)
    return _mask


def dilate_4_task0(sm_mask):
    sm_mask = np.array(sm_mask)
    preserve1 = [2, 3, 10, 5]
    mask1 = np.isin(sm_mask, preserve1)
    mask1 = _dilate(mask1, 7, 1)
    preserve2 = [3, 10]
    mask2 = np.isin(sm_mask, preserve2)
    mask2 = _dilate(mask2, 10, 3)
    preserve3 = [1]
    mask3 = np.isin(sm_mask, preserve3)
    mask3 = _dilate(mask3, 7, 2)
    mask = mask1 | mask2 | mask3
    return mask


class Dataset_custom(data.Dataset):
    mean = (0.5, 0.5, 0.5)
    std = (0.5, 0.5, 0.5)

    def get_img4clip(
        self,
        img,
        sm_mask,
        preserve,
        for_clip=True,
        add_semantic_head=False,
        mask_after_npisin=None,
        for_inpaint512=False,
    ):
        sm_mask = np.array(sm_mask)
        if mask_after_npisin is None:
            if self.task == 0 and 0:
                mask = dilate_4_task0(sm_mask)
            else:
                mask = np.isin(sm_mask, preserve)
                if self.task == 0 and 1 and for_inpaint512:
                    forehead_mask = get_forehead_mask(sm_mask)
                    mask = mask & ~forehead_mask
        else:
            mask = mask_after_npisin

        if isinstance(img, np.ndarray):
            img = Image.fromarray(img)
        if add_semantic_head:
            mask_before_colorSM = mask
            img, mask = add_colorSM(img, sm_mask, preserve, None)
        mask = mask_after_npisin__2__tensor(mask)

        if for_clip:
            image_tensor = get_tensor_clip()(img)
        else:
            image_tensor = get_tensor(mean=self.mean, std=self.std)(img)
        image_tensor = T.Resize([512, 512])(image_tensor)
        image_tensor = image_tensor * mask
        if for_clip:
            image_tensor = 255.0 * rearrange(un_norm_clip(image_tensor), "c h w -> h w c").cpu().numpy()
            _size = 224
        else:
            image_tensor = 255.0 * rearrange(un_norm(image_tensor), "c h w -> h w c").cpu().numpy()
            _size = 512
        
        image_tensor = albumentations.Resize(height=_size, width=_size)(image=image_tensor)
        image_tensor = Image.fromarray(image_tensor["image"].astype(np.uint8))
        if for_clip:
            image_tensor = get_tensor_clip()(image_tensor)
        else:
            image_tensor = get_tensor(mean=self.mean, std=self.std)(image_tensor)
            image_tensor = image_tensor * mask
        if add_semantic_head:
            mask = mask_after_npisin__2__tensor(mask_before_colorSM)
        return image_tensor, mask

    def __init__(
        self,
        state,
        task,
        paths_tgt,
        paths_ref,
        name="custom",
    ):
        if task == 0:
            USE_filter_mediapipe_fail_swap = 1
            USE_pts = 1
            READ_mediapipe_result_from_cache = 1
        elif task == 1:
            USE_filter_mediapipe_fail_swap = 0
            USE_pts = 0
            READ_mediapipe_result_from_cache = 1
        elif task == 2:
            USE_filter_mediapipe_fail_swap = 1
            USE_pts = 1
            READ_mediapipe_result_from_cache = 1
        elif task == 3:
            USE_filter_mediapipe_fail_swap = 0
            USE_pts = 1
            READ_mediapipe_result_from_cache = 1
        self.READ_mediapipe_result_from_cache = READ_mediapipe_result_from_cache

        assert state == "test"
        self.state = state
        self.image_size = 512
        self.kernel = np.ones((1, 1), np.uint8)
        self.name = name

        assert paths_tgt is not None and paths_ref is not None, "paths_tgt and paths_ref are required"
        assert len(paths_tgt) == len(paths_ref), "paths_tgt and paths_ref must be the same length"
        self.paths_tgt = list(paths_tgt)
        self.paths_ref = list(paths_ref)

        if READ_mediapipe_result_from_cache:
            self.mediapipe_Result_Cache = Mediapipe_Result_Cache()
        self.task = task

    def __getitem__(self, index):
        task = self.task
        path_tgt = self.paths_tgt[index]
        path_ref = self.paths_ref[index]


        img_tgt = Image.open(path_tgt).convert("RGB")
        img_tgt = resize_A(img_tgt, self.name)

        mask_path = path_img_2_path_mask(path_tgt)
        if self.task == 0:
            preserve = [1, 2, 3, 10, 5, 6, 7, 9]
            if 0:
                preserve = [1, 2, 3, 10, 5]
            sm_mask_tgt = Image.open(mask_path).convert("L")
            sm_mask_tgt = np.array(sm_mask_tgt)
            if 0:
                mask_tgt = dilate_4_task0(sm_mask_tgt)
            else:
                mask_tgt = np.isin(sm_mask_tgt, preserve)
                if self.task == 0 and 1:
                    forehead_mask = get_forehead_mask(sm_mask_tgt)
                    mask_tgt = mask_tgt & ~forehead_mask
        elif self.task == 1:
            preserve = [4]
            mask_tgt = path_img_2_mask(path_tgt, preserve)
        elif self.task == 3:
            preserve = [1, 2, 3, 10, 4, 5, 6, 7, 9]
            mask_tgt = path_img_2_mask(path_tgt, preserve)
        elif self.task == 2:
            preserve = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 20, 21]
            sm_mask_tgt = Image.open(mask_path).convert("L")
            sm_mask_tgt = np.array(sm_mask_tgt)
            mask_tgt = np.isin(sm_mask_tgt, preserve)

        converted_mask = np.zeros_like(mask_tgt)
        converted_mask[mask_tgt] = 255
        mask_tgt = Image.fromarray(converted_mask).convert("L")
        mask_tensor = 1 - get_tensor(normalize=False, toTensor=True)(mask_tgt)

        image_tensor = get_tensor(mean=self.mean, std=self.std)(img_tgt)
        image_tensor_resize = T.Resize([self.image_size, self.image_size])(image_tensor)
        mask_tensor_resize = T.Resize([self.image_size, self.image_size])(mask_tensor)

        if task == 2:
            inpaint_tensor_resize = image_tensor_resize
        else:
            inpaint_tensor_resize = image_tensor_resize * mask_tensor_resize
        if 1:
            mask_tensor_resize = 1 - mask_tensor_resize

        if 1:
            mask_path_ref = path_img_2_path_mask(path_ref)
            sm_mask_ref = Image.open(mask_path_ref).convert("L")
            sm_mask_ref = np.array(sm_mask_ref)
            img_ref = cv2.imread(str(path_ref))
            img_ref = cv2.cvtColor(img_ref, cv2.COLOR_BGR2RGB)
            img_ref = resize_A(img_ref, self.name)

        if task != 2:
            ref_image_tensor, ref_mask_tensor = self.get_img4clip(
                img_ref, sm_mask_ref, preserve, for_clip=True, add_semantic_head=0
            )
            if task == 3:
                ref_image_faceOnly_tensor, _ = self.get_img4clip(
                    img_ref,
                    sm_mask_ref,
                    [1, 2, 3, 10, 5, 6, 7, 9],
                    for_clip=False,
                    add_semantic_head=0,
                )
        else:
            ref_image_tensor = inpaint_tensor_resize

        ret = {
            "inpaint_image": inpaint_tensor_resize,
            "inpaint_mask": mask_tensor_resize,
            "ref_imgs": ref_image_tensor,
            "task": self.task,
        }

        if self.task == 0:
            ret["enInputs"] = {
                "face_ID-in": ref_image_tensor,
                "face-clip-in": ref_image_tensor,
            }
        elif self.task == 1:
            ret["enInputs"] = {
                "hair-clip-in": ref_image_tensor,
            }
        elif self.task == 2:
            tgt_nonBg_tensor, _ = self.get_img4clip(img_tgt, sm_mask_tgt, preserve)
            ret["enInputs"] = {
                "face_ID-in": tgt_nonBg_tensor,
                "head-clip-in": tgt_nonBg_tensor,
            }
        elif self.task == 3:
            ret["enInputs"] = {
                "face_ID-in": ref_image_faceOnly_tensor,
                "head-clip-in": ref_image_tensor,
            }

        if (REFNET.ENABLE and REFNET.task2layerNum[task] > 0) or CH14:
            if task != 2:
                ref_imgs_4unet, ref_mask_4unet = self.get_img4clip(
                    img_ref, sm_mask_ref, preserve, for_clip=False, add_semantic_head=0
                )
            else:
                ref_imgs_4unet, ref_mask_4unet = self.get_img4clip(
                    img_tgt,
                    sm_mask_tgt,
                    "any",
                    for_clip=False,
                    add_semantic_head=0,
                    mask_after_npisin=np.ones_like(sm_mask_tgt).astype(bool),
                )
            ref_imgs_4unet = T.Resize([self.image_size, self.image_size])(ref_imgs_4unet)
            ref_mask_512 = T.Resize([self.image_size, self.image_size])(ref_mask_4unet)
            ret["ref_imgs_4unet"] = ref_imgs_4unet
            ret["ref_mask_512"] = ref_mask_512

        if self.READ_mediapipe_result_from_cache:
            if self.state == "test":
                if task == 2:
                    _p_lmk = path_ref
                else:
                    _p_lmk = path_tgt
            else:
                _p_lmk = path_tgt
            ret["mediapipe_lmkAll"] = self.mediapipe_Result_Cache.get(_p_lmk)
            if ret["mediapipe_lmkAll"] is None:
                raise RuntimeError(
                    f"Missing Mediapipe cache for input image: {_p_lmk}. "
                    "Precompute landmarks and ensure cache exists before inference."
                )

        if self.state == "test":
            prior_image_tensor = "None"
            out_stem = f"{Path(path_tgt).stem}-{Path(path_ref).stem}"
            if task == 2:
                ref512, _ = self.get_img4clip(
                    img_ref, sm_mask_ref, preserve, for_clip=False, add_semantic_head=0
                )
                ref512 = T.Resize([self.image_size, self.image_size])(ref512)
                ret["ref512"] = ref512
            ret = (image_tensor_resize, prior_image_tensor, ret, out_stem)
        return ret

    def __len__(self):
        return len(self.paths_tgt)