import cv2
import numpy as np
import torch
import threading
from torchvision import transforms
from clip.clipseg import CLIPDensePredT
import numpy as np

from roop.typing import Frame

THREAD_LOCK_CLIP = threading.Lock()


class Mask_Clip2Seg:
    plugin_options: dict = None
    model_clip = None

    processorname = "clip2seg"
    type = "mask"

    def Initialize(self, plugin_options: dict):
        if self.plugin_options is not None:
            if self.plugin_options["devicename"] != plugin_options["devicename"]:
                self.Release()

        self.plugin_options = plugin_options
        if self.model_clip is None:
            self.model_clip = CLIPDensePredT(
                version="ViT-B/16", reduce_dim=64, complex_trans_conv=True
            )
            self.model_clip.eval()
            self.model_clip.load_state_dict(
                torch.load(
                    "models/CLIP/rd64-uni-refined.pth", map_location=torch.device("cpu")
                ),
                strict=False,
            )

        device = torch.device(self.plugin_options["devicename"])
        self.model_clip.to(device)

    def Run(self, img1, keywords: str) -> Frame:
        if keywords is None or len(keywords) < 1 or img1 is None:
            return img1

        source_image_small = cv2.resize(img1, (256, 256))

        img_mask = np.full(
            (source_image_small.shape[0], source_image_small.shape[1]),
            0,
            dtype=np.float32,
        )
        mask_border = 1
        l = 0
        t = 0
        r = 1
        b = 1

        mask_blur = 5
        clip_blur = 5

        img_mask = cv2.rectangle(
            img_mask,
            (mask_border + int(l), mask_border + int(t)),
            (256 - mask_border - int(r), 256 - mask_border - int(b)),
            (255, 255, 255),
            -1,
        )
        img_mask = cv2.GaussianBlur(img_mask, (mask_blur * 2 + 1, mask_blur * 2 + 1), 0)
        img_mask /= 255

        input_image = source_image_small

        transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                ),
                transforms.Resize((256, 256)),
            ]
        )
        img = transform(input_image).unsqueeze(0)

        thresh = 0.5
        prompts = keywords.split(",")
        with THREAD_LOCK_CLIP:
            with torch.no_grad():
                preds = self.model_clip(img.repeat(len(prompts), 1, 1, 1), prompts)[0]
        clip_mask = torch.sigmoid(preds[0][0])
        for i in range(len(prompts) - 1):
            clip_mask += torch.sigmoid(preds[i + 1][0])

        clip_mask = clip_mask.data.cpu().numpy()
        np.clip(clip_mask, 0, 1)

        clip_mask[clip_mask > thresh] = 1.0
        clip_mask[clip_mask <= thresh] = 0.0
        kernel = np.ones((5, 5), np.float32)
        clip_mask = cv2.dilate(clip_mask, kernel, iterations=1)
        clip_mask = cv2.GaussianBlur(
            clip_mask, (clip_blur * 2 + 1, clip_blur * 2 + 1), 0
        )

        img_mask *= clip_mask
        img_mask[img_mask < 0.0] = 0.0
        return img_mask

    def Release(self):
        self.model_clip = None