import os
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from backend.patcher.unet import UnetPatcher
    from backend.patcher.vae import VAE
    from modules.processing import StableDiffusionProcessing

import cv2
import einops
import numpy as np
import torch
import torch.nn.functional as F
import yaml
from omegaconf import OmegaConf

from modules.modelloader import load_file_from_url
from modules_forge.shared import add_supported_preprocessor, preprocessor_dir
from modules_forge.supported_preprocessor import Preprocessor, PreprocessorParameter
from modules_forge.utils import resize_image_with_pad


class PreprocessorInpaint(Preprocessor):
    def __init__(self):
        super().__init__()
        self.name = "inpaint_global_harmonious"
        self.tags = ["Inpaint"]
        self.model_filename_filters = ["inpaint"]
        self.slider_resolution = PreprocessorParameter(visible=False)
        self.fill_mask_with_one_when_resize_and_fill = True
        self.expand_mask_when_resize_and_fill = True

    @torch.inference_mode()
    def process_before_every_sampling(self, process: "StableDiffusionProcessing", cond: torch.Tensor, mask: torch.Tensor, *args, **kwargs):
        mask = mask.round()
        mixed_cond = cond * (1.0 - mask) - mask
        return mixed_cond, None


class PreprocessorInpaintOnly(PreprocessorInpaint):
    def __init__(self):
        super().__init__()
        self.name = "inpaint_only"
        self.image = None
        self.mask = None
        self.latent = None

    @torch.inference_mode()
    def process_before_every_sampling(self, process: "StableDiffusionProcessing", cond: torch.Tensor, mask: torch.Tensor, *args, **kwargs):
        mask = mask.round()
        self.image = cond
        self.mask = mask

        vae: "VAE" = process.sd_model.forge_objects.vae

        latent_image = vae.encode(self.image.movedim(1, -1))
        latent_image = process.sd_model.forge_objects.vae.first_stage_model.process_in(latent_image)

        *_, H, W = latent_image.shape

        latent_mask = self.mask
        latent_mask = torch.nn.functional.interpolate(latent_mask, size=(H * 8, W * 8), mode="bilinear").round()
        latent_mask = torch.nn.functional.max_pool2d(latent_mask, (8, 8)).round().to(latent_image)

        unet: "UnetPatcher" = process.sd_model.forge_objects.unet.clone()

        def pre_cfg(model, c, uc, x, timestep, model_options):
            noisy_latent = latent_image.to(x) + timestep[:, None, None, None].to(x) * torch.randn_like(latent_image).to(x)
            x = x * latent_mask.to(x) + noisy_latent.to(x) * (1.0 - latent_mask.to(x))
            return model, c, uc, x, timestep, model_options

        def post_cfg(args):
            denoised = args["denoised"]
            denoised = denoised * latent_mask.to(denoised) + latent_image.to(denoised) * (1.0 - latent_mask.to(denoised))
            return denoised

        unet.add_sampler_pre_cfg_function(pre_cfg)
        unet.set_model_sampler_post_cfg_function(post_cfg)

        process.sd_model.forge_objects.unet = unet

        self.latent = latent_image

        mixed_cond = cond * (1.0 - mask) - mask

        return mixed_cond, None

    def process_after_every_sampling(self, process, params, *args, **kwargs):
        a1111_batch_result = args[0]
        new_results = []

        for img in a1111_batch_result.images:
            sigma = 7
            mask = self.mask[0, 0].detach().cpu().numpy().astype(np.float32)
            mask = cv2.dilate(mask, np.ones((sigma, sigma), dtype=np.uint8))
            mask = cv2.blur(mask, (sigma, sigma))[None]
            mask = torch.from_numpy(np.ascontiguousarray(mask).copy()).to(img).clip(0, 1)
            raw = self.image[0].to(img).clip(0, 1)
            img = img.clip(0, 1)

            if img.shape != mask.shape:  # Interrupted
                img = F.interpolate(img.unsqueeze(0), size=(mask.shape[1], mask.shape[2]), mode="nearest").squeeze(0)

            new_results.append(raw * (1.0 - mask) + img * mask)

        a1111_batch_result.images = new_results


class PreprocessorInpaintLama(PreprocessorInpaintOnly):
    def __init__(self):
        super().__init__()
        self.name = "inpaint_only+lama"

    def load_model(self):
        remote_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetLama.pth"
        model_path = load_file_from_url(remote_model_path, model_dir=preprocessor_dir)

        config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lama_config.yaml")
        with open(config_path, "r") as f:
            cfg = yaml.safe_load(f)

        cfg = OmegaConf.create(cfg)
        cfg.training_model.predict_only = True
        cfg.visualizer.kind = "noop"

        from lama import load_checkpoint

        model = load_checkpoint(cfg, os.path.abspath(model_path), map_location="cpu", strict=False)

        self.setup_model_patcher(model)

    def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, input_mask=None, **kwargs):
        if input_mask is None:
            return input_image

        H, W, _ = input_image.shape
        raw_color = input_image.copy()
        raw_mask = input_mask.copy()

        input_image, remove_pad = resize_image_with_pad(input_image, 256)
        input_mask, remove_pad = resize_image_with_pad(input_mask, 256)
        input_mask = input_mask[..., :1]

        self.load_model()

        self.move_all_model_patchers_to_gpu()

        color = np.ascontiguousarray(input_image).astype(np.float32) / 255.0
        mask = np.ascontiguousarray(input_mask).astype(np.float32) / 255.0
        with torch.inference_mode():
            color = self.send_tensor_to_model_device(torch.from_numpy(color))
            mask = self.send_tensor_to_model_device(torch.from_numpy(mask))
            mask = (mask > 0.5).float()
            color = color * (1 - mask)
            image_feed = torch.cat([color, mask], dim=2)
            image_feed = einops.rearrange(image_feed, "h w c -> 1 c h w")
            prd_color = self.model_patcher.model(image_feed)[0]
            prd_color = einops.rearrange(prd_color, "c h w -> h w c")
            prd_color = prd_color * mask + color * (1 - mask)
            prd_color *= 255.0
            prd_color = prd_color.detach().cpu().numpy().clip(0, 255).astype(np.uint8)

        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)

        return fin_color

    @torch.inference_mode()
    def process_before_every_sampling(self, process: "StableDiffusionProcessing", cond: torch.Tensor, mask: torch.Tensor, *args, **kwargs):
        cond, mask = super().process_before_every_sampling(process, cond, mask, *args, **kwargs)
        sigma_max = process.sd_model.forge_objects.unet.model.predictor.sigma_max
        original_noise = kwargs["noise"]
        process.modified_noise = original_noise + self.latent.to(original_noise) / sigma_max.to(original_noise)
        return cond, mask


class PreprocessorInpaintNoobAIXL(PreprocessorInpaint):
    def __init__(self):
        super().__init__()
        self.name = "inpaint_noobai"
        self.tags = ["Inpaint"]
        self.model_filename_filters = ["inpaint", "noobai"]

    def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, input_mask=None, **kwargs):
        if input_mask is None:
            return input_image

        if not isinstance(input_image, np.ndarray):
            input_image = np.array(input_image)
        if not isinstance(input_mask, np.ndarray):
            input_mask = np.array(input_mask)

        mask = input_mask.astype(np.float32) / 255.0
        mask = (mask > 0.5).astype(bool)

        if mask.ndim == 2:
            mask = np.expand_dims(mask, axis=-1)
        if mask.shape[-1] == 1:
            mask = np.repeat(mask, 3, axis=-1)

        result = input_image.copy()
        result[mask] = 0.0

        return result

    @torch.inference_mode()
    def process_before_every_sampling(self, process: "StableDiffusionProcessing", cond: torch.Tensor, mask: torch.Tensor, *args, **kwargs):
        if process.denoising_strength < 0.9:
            print("High Denoising Strength is Recommended!")

        mask = mask.round()
        mixed_cond = cond * (1.0 - mask)

        return mixed_cond, None


add_supported_preprocessor(PreprocessorInpaint())

add_supported_preprocessor(PreprocessorInpaintOnly())

add_supported_preprocessor(PreprocessorInpaintLama())

add_supported_preprocessor(PreprocessorInpaintNoobAIXL())
