# 1st Edit by. https://github.com/shiimizu/ComfyUI-TiledDiffusion
# 2nd Edit by. Forge Official
# 3rd Edit by. Haoming02
# - Based on: https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111

from enum import Enum
from typing import Final, Optional, Union

import numpy as np
import torch
from numpy import exp, pi, sqrt
from torch import Tensor

from backend import memory_management
from backend.misc.image_resize import adaptive_resize
from backend.patcher.base import ModelPatcher
from backend.patcher.controlnet import ControlNet, T2IAdapter

_dev: Final[torch.device] = memory_management.get_torch_device()

opt_f: Optional[int] = None


class BlendMode(Enum):
    FOREGROUND = "Foreground"
    BACKGROUND = "Background"


class BBox:

    def __init__(self, x: int, y: int, w: int, h: int):
        self.x = x
        self.y = y
        self.w = w
        self.h = h
        self.box = [x, y, x + w, y + h]
        self.slicer = slice(None), slice(None), slice(y, y + h), slice(x, x + w)

    def __getitem__(self, idx: int) -> int:
        return self.box[idx]


def ceildiv(big, small):
    return -(big // -small)


def split_bboxes(w: int, h: int, tile_w: int, tile_h: int, overlap: int = 16, init_weight: Union[Tensor, float] = 1.0) -> tuple[list[BBox], Tensor]:
    cols = ceildiv((w - overlap), (tile_w - overlap))
    rows = ceildiv((h - overlap), (tile_h - overlap))
    dx = (w - tile_w) / (cols - 1) if cols > 1 else 0
    dy = (h - tile_h) / (rows - 1) if rows > 1 else 0

    bbox_list: list[BBox] = []
    weight = torch.zeros((1, 1, h, w), device=_dev, dtype=torch.float32)
    for row in range(rows):
        y = min(int(row * dy), h - tile_h)
        for col in range(cols):
            x = min(int(col * dx), w - tile_w)

            bbox = BBox(x, y, tile_w, tile_h)
            bbox_list.append(bbox)
            weight[bbox.slicer] += init_weight

    return bbox_list, weight


class AbstractDiffusion:
    def __init__(self):
        self.method = self.__class__.__name__
        self.pbar = None

        self.w: int = 0
        self.h: int = 0
        self.tile_width: int = None
        self.tile_height: int = None
        self.tile_overlap: int = None
        self.tile_batch_size: int = None
        self.x_buffer: Tensor = None
        self._weights: Tensor = None
        self._init_grid_bbox = None
        self._init_done = None
        self.step_count = 0
        self.inner_loop_count = 0
        self.kdiff_step = -1
        self.enable_grid_bbox: bool = False
        self.tile_w: int = None
        self.tile_h: int = None
        self.tile_bs: int = None
        self.num_tiles: int = None
        self.num_batches: int = None
        self.batched_bboxes: list[list[BBox]] = []
        self.enable_custom_bbox: bool = False
        self.custom_bboxes: list[BBox] = []
        self.enable_controlnet: bool = False
        self.control_tensor_batch_dict = {}
        self.control_tensor_batch: list[list[Tensor]] = [[]]
        self.control_params: dict[tuple, list[list[Tensor]]] = {}
        self.control_tensor_cpu: bool = None
        self.control_tensor_custom: list[list[Tensor]] = []

        self.draw_background: bool = True
        self.control_tensor_cpu = False
        self.weights = None

    def reset(self):
        tile_width = self.tile_width
        tile_height = self.tile_height
        tile_overlap = self.tile_overlap
        tile_batch_size = self.tile_batch_size
        self.__init__()
        self.tile_width = tile_width
        self.tile_height = tile_height
        self.tile_overlap = tile_overlap
        self.tile_batch_size = tile_batch_size

    def repeat_tensor(self, x: Tensor, n: int, concat=False, concat_to=0) -> Tensor:
        if n == 1:
            return x
        B = x.shape[0]
        r_dims = len(x.shape) - 1
        if B == 1:
            shape = [n] + [-1] * r_dims
            return x.expand(shape)
        else:
            if concat:
                return torch.cat([x for _ in range(n)], dim=0)[:concat_to]
            shape = [n] + [1] * r_dims
            return x.repeat(shape)

    def update_pbar(self):
        if self.pbar.n >= self.pbar.total:
            self.pbar.close()
        else:
            sampling_step = 20
            if self.step_count == sampling_step:
                self.inner_loop_count += 1
                if self.inner_loop_count < self.total_bboxes:
                    self.pbar.update()
            else:
                self.step_count = sampling_step
                self.inner_loop_count = 0

    def reset_buffer(self, x_in: Tensor):
        if self.x_buffer is None or self.x_buffer.shape != x_in.shape:
            self.x_buffer = torch.zeros_like(x_in, device=x_in.device, dtype=x_in.dtype)
        else:
            self.x_buffer.zero_()

    def init_grid_bbox(self, tile_w: int, tile_h: int, overlap: int, tile_bs: int):
        self.weights = torch.zeros((1, 1, self.h, self.w), device=_dev, dtype=torch.float32)
        self.enable_grid_bbox = True

        self.tile_w = min(tile_w, self.w)
        self.tile_h = min(tile_h, self.h)
        overlap = max(0, min(overlap, min(tile_w, tile_h) - 4))
        bboxes, weights = split_bboxes(self.w, self.h, self.tile_w, self.tile_h, overlap, self.get_tile_weights())
        self.weights += weights
        self.num_tiles = len(bboxes)
        self.num_batches = ceildiv(self.num_tiles, tile_bs)
        self.tile_bs = ceildiv(len(bboxes), self.num_batches)
        self.batched_bboxes = [bboxes[i * self.tile_bs : (i + 1) * self.tile_bs] for i in range(self.num_batches)]

    def get_tile_weights(self) -> Union[Tensor, float]:
        return 1.0

    def init_noise_inverse(self, steps: int, retouch: float, get_cache_callback, set_cache_callback, renoise_strength: float, renoise_kernel: int):
        self.noise_inverse_enabled = True
        self.noise_inverse_steps = steps
        self.noise_inverse_retouch = float(retouch)
        self.noise_inverse_renoise_strength = float(renoise_strength)
        self.noise_inverse_renoise_kernel = int(renoise_kernel)
        self.noise_inverse_set_cache = set_cache_callback
        self.noise_inverse_get_cache = get_cache_callback

    def init_done(self):
        self.total_bboxes = 0
        if self.enable_grid_bbox:
            self.total_bboxes += self.num_batches
        if self.enable_custom_bbox:
            self.total_bboxes += len(self.custom_bboxes)
        assert self.total_bboxes > 0, "Nothing to paint! No background to draw and no custom bboxes were provided."

    def prepare_controlnet_tensors(self, refresh: bool = False, tensor=None):
        if not refresh:
            if self.control_tensor_batch is not None or self.control_params is not None:
                return
        tensors = [tensor]
        self.org_control_tensor_batch = tensors
        self.control_tensor_batch = []
        for i in range(len(tensors)):
            control_tile_list = []
            control_tensor = tensors[i]
            for bboxes in self.batched_bboxes:
                single_batch_tensors = []
                for bbox in bboxes:
                    if len(control_tensor.shape) == 3:
                        control_tensor.unsqueeze_(0)
                    control_tile = control_tensor[:, :, bbox[1] * opt_f : bbox[3] * opt_f, bbox[0] * opt_f : bbox[2] * opt_f]
                    single_batch_tensors.append(control_tile)
                control_tile = torch.cat(single_batch_tensors, dim=0)
                if self.control_tensor_cpu:
                    control_tile = control_tile.cpu()
                control_tile_list.append(control_tile)
            self.control_tensor_batch.append(control_tile_list)

            if len(self.custom_bboxes) > 0:
                custom_control_tile_list = []
                for bbox in self.custom_bboxes:
                    if len(control_tensor.shape) == 3:
                        control_tensor.unsqueeze_(0)
                    control_tile = control_tensor[:, :, bbox[1] * opt_f : bbox[3] * opt_f, bbox[0] * opt_f : bbox[2] * opt_f]
                    if self.control_tensor_cpu:
                        control_tile = control_tile.cpu()
                    custom_control_tile_list.append(control_tile)
                self.control_tensor_custom.append(custom_control_tile_list)

    def switch_controlnet_tensors(self, batch_id: int, x_batch_size: int, tile_batch_size: int, is_denoise=False):
        if self.control_tensor_batch is None:
            return
        for param_id in range(len(self.control_tensor_batch)):
            control_tile = self.control_tensor_batch[param_id][batch_id]
            if x_batch_size > 1:
                all_control_tile = []
                for i in range(tile_batch_size):
                    this_control_tile = [control_tile[i].unsqueeze(0)] * x_batch_size
                    all_control_tile.append(torch.cat(this_control_tile, dim=0))
                control_tile = torch.cat(all_control_tile, dim=0)
                self.control_tensor_batch[param_id][batch_id] = control_tile

    def process_controlnet(self, x_shape, x_dtype, c_in: dict, cond_or_uncond: list, bboxes, batch_size: int, batch_id: int):
        control: ControlNet = c_in["control_model"]
        param_id = -1
        tuple_key = tuple(cond_or_uncond) + tuple(x_shape)
        while control is not None:
            param_id += 1
            PH, PW = self.h * 8, self.w * 8

            if self.control_params.get(tuple_key, None) is None:
                self.control_params[tuple_key] = [[None]]
                val = self.control_params[tuple_key]
                if param_id + 1 >= len(val):
                    val.extend([[None] for _ in range(param_id + 1)])
                if len(self.batched_bboxes) >= len(val[param_id]):
                    val[param_id].extend([[None] for _ in range(len(self.batched_bboxes))])
            if self.refresh or control.cond_hint is None or not isinstance(self.control_params[tuple_key][param_id][batch_id], Tensor):
                dtype = getattr(control, "manual_cast_dtype", None)
                if dtype is None:
                    dtype = getattr(getattr(control, "control_model", None), "dtype", None)
                if dtype is None:
                    dtype = x_dtype
                if isinstance(control, T2IAdapter):
                    width, height = control.scale_image_to(PW, PH)
                    control.cond_hint = adaptive_resize(control.cond_hint_original, width, height, "nearest-exact", "center").float().to(control.device)
                    if control.channels_in == 1 and control.cond_hint.shape[1] > 1:
                        control.cond_hint = torch.mean(control.cond_hint, 1, keepdim=True)
                elif control.__class__.__name__ == "ControlLLLiteAdvanced":
                    if control.sub_idxs is not None and control.cond_hint_original.shape[0] >= control.full_latent_length:
                        control.cond_hint = adaptive_resize(control.cond_hint_original[control.sub_idxs], PW, PH, "nearest-exact", "center").to(dtype=dtype, device=control.device)
                    else:
                        if (PH, PW) == (control.cond_hint_original.shape[-2], control.cond_hint_original.shape[-1]):
                            control.cond_hint = control.cond_hint_original.clone().to(dtype=dtype, device=control.device)
                        else:
                            control.cond_hint = adaptive_resize(control.cond_hint_original, PW, PH, "nearest-exact", "center").to(dtype=dtype, device=control.device)
                else:
                    if (PH, PW) == (control.cond_hint_original.shape[-2], control.cond_hint_original.shape[-1]):
                        control.cond_hint = control.cond_hint_original.clone().to(dtype=dtype, device=control.device)
                    else:
                        control.cond_hint = adaptive_resize(control.cond_hint_original, PW, PH, "nearest-exact", "center").to(dtype=dtype, device=control.device)
                cond_hint_pre_tile = control.cond_hint
                if control.cond_hint.shape[0] < batch_size:
                    cond_hint_pre_tile = self.repeat_tensor(control.cond_hint, ceildiv(batch_size, control.cond_hint.shape[0]))[:batch_size]
                cns = [cond_hint_pre_tile[:, :, bbox[1] * opt_f : bbox[3] * opt_f, bbox[0] * opt_f : bbox[2] * opt_f] for bbox in bboxes]
                control.cond_hint = torch.cat(cns, dim=0)
                self.control_params[tuple_key][param_id][batch_id] = control.cond_hint
            else:
                control.cond_hint = self.control_params[tuple_key][param_id][batch_id]
            control = control.previous_controlnet


def gaussian_weights(tile_w: int, tile_h: int) -> Tensor:
    f = lambda x, midpoint, var=0.01: exp(-(x - midpoint) * (x - midpoint) / (tile_w * tile_w) / (2 * var)) / sqrt(2 * pi * var)
    x_probs = [f(x, (tile_w - 1) / 2) for x in range(tile_w)]
    y_probs = [f(y, tile_h / 2) for y in range(tile_h)]

    w = np.outer(y_probs, x_probs)
    return torch.from_numpy(w).to(_dev, dtype=torch.float32)


class MultiDiffusion(AbstractDiffusion):

    @torch.no_grad()
    def __call__(self, model_function, args: dict):
        x_in: Tensor = args["input"]
        t_in: Tensor = args["timestep"]
        c_in: dict = args["c"]
        cond_or_uncond: list = args["cond_or_uncond"]
        c_crossattn: Tensor = c_in["c_crossattn"]

        is_5d = x_in.ndim == 5
        if is_5d:
            assert x_in.shape[2] == 1
            x_in = x_in.squeeze(2)

        N, C, H, W = x_in.shape
        self.refresh = False
        if self.weights is None or self.h != H or self.w != W:
            self.h, self.w = H, W
            self.refresh = True
            self.init_grid_bbox(self.tile_width, self.tile_height, self.tile_overlap, self.tile_batch_size)
            self.init_done()
        self.h, self.w = H, W
        self.reset_buffer(x_in)
        if self.draw_background:
            for batch_id, bboxes in enumerate(self.batched_bboxes):
                x_tile = torch.cat([x_in[bbox.slicer] for bbox in bboxes], dim=0)
                n_rep = len(bboxes)
                ts_tile = self.repeat_tensor(t_in, n_rep)
                cond_tile = self.repeat_tensor(c_crossattn, n_rep)
                c_tile = c_in.copy()
                c_tile["c_crossattn"] = cond_tile
                if "time_context" in c_in:
                    c_tile["time_context"] = self.repeat_tensor(c_in["time_context"], n_rep)
                for key in c_tile:
                    if key in ["y", "c_concat"]:
                        icond = c_tile[key]
                        if icond.ndim == 5:
                            assert icond.shape[2] == 1
                            icond = icond.squeeze(2)

                        if icond.shape[2:] == (self.h, self.w):
                            c_tile[key] = torch.cat([icond[bbox.slicer] for bbox in bboxes])
                        else:
                            c_tile[key] = self.repeat_tensor(icond, n_rep)
                if "control" in c_in:
                    self.process_controlnet(x_tile.shape, x_tile.dtype, c_in, cond_or_uncond, bboxes, N, batch_id)
                    c_tile["control"] = c_in["control_model"].get_control(x_tile, ts_tile, c_tile, len(cond_or_uncond))

                if is_5d:
                    x_tile = x_tile.unsqueeze(2)
                    for key in ["y", "c_concat"]:
                        if key in c_tile and c_tile[key].ndim == 4:
                            c_tile[key] = c_tile[key].unsqueeze(2)

                    control = c_tile.get("control")
                    while control is not None:
                        if hasattr(control, "cond_hint") and control.cond_hint is not None and control.cond_hint.ndim == 4:
                            control.cond_hint = control.cond_hint.unsqueeze(2)
                        control = control.previous_controlnet

                x_tile_out = model_function(x_tile, ts_tile, **c_tile)

                if is_5d:
                    x_tile_out = x_tile_out.squeeze(2)

                for i, bbox in enumerate(bboxes):
                    self.x_buffer[bbox.slicer] += x_tile_out[i * N : (i + 1) * N, :, :, :]
                del x_tile_out, x_tile, ts_tile, c_tile
        x_out = torch.where(self.weights > 1, self.x_buffer / self.weights, self.x_buffer)

        if is_5d:
            x_out = x_out.unsqueeze(2)

        return x_out


class MixtureOfDiffusers(AbstractDiffusion):
    """
    Mixture-of-Diffusers Implementation
    https://github.com/albarji/mixture-of-diffusers
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.custom_weights: list[Tensor] = []
        self.get_weight = gaussian_weights

    def init_done(self):
        super().init_done()
        self.rescale_factor = 1 / self.weights
        for bbox_id, bbox in enumerate(self.custom_bboxes):
            if bbox.blend_mode == BlendMode.BACKGROUND:
                self.custom_weights[bbox_id] *= self.rescale_factor[bbox.slicer]

    def get_tile_weights(self) -> Tensor:
        self.tile_weights = self.get_weight(self.tile_w, self.tile_h)
        return self.tile_weights

    @torch.no_grad()
    def __call__(self, model_function, args: dict):
        x_in: Tensor = args["input"]
        t_in: Tensor = args["timestep"]
        c_in: dict = args["c"]
        cond_or_uncond: list = args["cond_or_uncond"]
        c_crossattn: Tensor = c_in["c_crossattn"]

        is_5d = x_in.ndim == 5
        if is_5d:
            assert x_in.shape[2] == 1
            x_in = x_in.squeeze(2)

        N, C, H, W = x_in.shape

        self.refresh = False
        if self.weights is None or self.h != H or self.w != W:
            self.h, self.w = H, W
            self.refresh = True
            self.init_grid_bbox(self.tile_width, self.tile_height, self.tile_overlap, self.tile_batch_size)
            self.init_done()
        self.h, self.w = H, W
        self.reset_buffer(x_in)
        if self.draw_background:
            for batch_id, bboxes in enumerate(self.batched_bboxes):
                x_tile_list = []
                t_tile_list = []
                icond_map = {}
                for bbox in bboxes:
                    x_tile_list.append(x_in[bbox.slicer])
                    t_tile_list.append(t_in)
                    if isinstance(c_in, dict):
                        for key in ["y", "c_concat"]:
                            if key in c_in:
                                icond = c_in[key]
                                if icond.ndim == 5:
                                    assert icond.shape[2] == 1
                                    icond = icond.squeeze(2)

                                if icond.shape[2:] == (self.h, self.w):
                                    icond = icond[bbox.slicer]
                                if icond_map.get(key, None) is None:
                                    icond_map[key] = []
                                icond_map[key].append(icond)
                    else:
                        print(">> [WARN] not supported, make an issue on github!!")
                n_rep = len(bboxes)
                x_tile = torch.cat(x_tile_list, dim=0)
                t_tile = self.repeat_tensor(t_in, n_rep)
                tcond_tile = self.repeat_tensor(c_crossattn, n_rep)
                c_tile = c_in.copy()
                c_tile["c_crossattn"] = tcond_tile
                if "time_context" in c_in:
                    c_tile["time_context"] = self.repeat_tensor(c_in["time_context"], n_rep)
                for key in c_tile:
                    if key in ["y", "c_concat"]:
                        icond_tile = torch.cat(icond_map[key], dim=0)
                        c_tile[key] = icond_tile
                if "control" in c_in:
                    self.process_controlnet(x_tile.shape, x_tile.dtype, c_in, cond_or_uncond, bboxes, N, batch_id)
                    c_tile["control"] = c_in["control_model"].get_control(x_tile, t_tile, c_tile, len(cond_or_uncond))

                if is_5d:
                    x_tile = x_tile.unsqueeze(2)
                    for key in ["y", "c_concat"]:
                        if key in c_tile and c_tile[key].ndim == 4:
                            c_tile[key] = c_tile[key].unsqueeze(2)

                    control = c_tile.get("control")
                    while control is not None:
                        if hasattr(control, "cond_hint") and control.cond_hint is not None and control.cond_hint.ndim == 4:
                            control.cond_hint = control.cond_hint.unsqueeze(2)
                        control = control.previous_controlnet

                x_tile_out = model_function(x_tile, t_tile, **c_tile)

                if is_5d:
                    x_tile_out = x_tile_out.squeeze(2)

                for i, bbox in enumerate(bboxes):
                    w = self.tile_weights * self.rescale_factor[bbox.slicer]
                    self.x_buffer[bbox.slicer] += x_tile_out[i * N : (i + 1) * N, :, :, :] * w
                del x_tile_out, x_tile, t_tile, c_tile
        x_out = self.x_buffer

        if is_5d:
            x_out = x_out.unsqueeze(2)

        return x_out


class TiledDiffusion:

    @staticmethod
    def apply(model: ModelPatcher, method: str, tile_width: int, tile_height: int, tile_overlap: int, tile_batch_size: int):
        match method:
            case "MultiDiffusion":
                implement = MultiDiffusion()
            case "Mixture of Diffusers":
                implement = MixtureOfDiffusers()
            case _:
                raise SystemError

        from modules import processing

        global opt_f
        opt_f = processing.opt_f

        implement.tile_width = tile_width // opt_f
        implement.tile_height = tile_height // opt_f
        implement.tile_overlap = tile_overlap // opt_f
        implement.tile_batch_size = tile_batch_size

        model = model.clone()
        model.set_model_unet_function_wrapper(implement)
        model.model_options["tiled_diffusion"] = True
        return model
