import cv2
import numpy as np

from ..supported_preprocessor import Preprocessor, PreprocessorParameter
from ..utils import resize_image_with_pad


class PreprocessorLamaInpaint(Preprocessor):
    def __init__(self):
        super().__init__(name="inpaint_only+lama")
        self.tags = ["Inpaint"]
        self.returns_image = True
        self.model = None
        self.slider_resolution = PreprocessorParameter(visible=False)

    def __call__(
        self,
        input_image,
        resolution,
        slider_1=None,
        slider_2=None,
        slider_3=None,
        **kwargs
    ):
        img = input_image
        H, W, C = img.shape
        raw_color = img[:, :, 0:3].copy()
        raw_mask = img[:, :, 3:4].copy()

        res = 256  # Always use 256 since lama is trained on 256

        img_res, remove_pad = resize_image_with_pad(img, res)

        if self.model is None:
            from annotator.lama import LamaInpainting

            self.model = LamaInpainting()
        # applied auto inversion
        prd_color = self.model(img_res)
        self.model.unload_model()

        prd_color = remove_pad(prd_color)
        prd_color = cv2.resize(prd_color, (W, H))

        alpha = raw_mask.astype(np.float32) / 255.0
        fin_color = prd_color.astype(np.float32) * alpha + raw_color.astype(
            np.float32
        ) * (1 - alpha)
        fin_color = fin_color.clip(0, 255).astype(np.uint8)

        result = np.concatenate([fin_color, raw_mask], axis=2)
        return result


Preprocessor.add_supported_preprocessor(PreprocessorLamaInpaint())