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

import comfy.model_management as model_management
from comfy.utils import common_upscale

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

LATENT_SCALE_FACTOR = 8


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: int, small: int) -> int:
    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


def _resize_hint(tensor: Tensor, width: int, height: int, like: Tensor) -> Tensor:
    resized = common_upscale(tensor, width, height, "nearest-exact", "center")
    return resized.to(device=like.device, dtype=like.dtype)


class AbstractDiffusion:
    def __init__(self):
        self.method = self.__class__.__name__
        self.w = 0
        self.h = 0
        self.tile_width: Optional[int] = None
        self.tile_height: Optional[int] = None
        self.tile_overlap: Optional[int] = None
        self.tile_batch_size: Optional[int] = None
        self.x_buffer: Optional[Tensor] = None
        self.enable_grid_bbox = False
        self.tile_w: Optional[int] = None
        self.tile_h: Optional[int] = None
        self.tile_bs: Optional[int] = None
        self.num_tiles: Optional[int] = None
        self.num_batches: Optional[int] = None
        self.batched_bboxes: list[list[BBox]] = []
        self.enable_custom_bbox = False
        self.custom_bboxes: list[BBox] = []
        self.draw_background = True
        self.weights: Optional[Tensor] = None
        self.control_params: dict[tuple, list[list[Optional[Tensor]]]] = {}

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

    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_done(self):
        total_bboxes = 0
        if self.enable_grid_bbox:
            total_bboxes += self.num_batches or 0
        if self.enable_custom_bbox:
            total_bboxes += len(self.custom_bboxes)
        if total_bboxes <= 0:
            raise AssertionError("Nothing to paint! No background to draw and no custom bboxes were provided.")

    def process_controlnet(self, x_like: Tensor, c_in: dict, cond_or_uncond: list, bboxes, batch_size: int, batch_id: int):
        control = c_in.get("control")
        param_id = -1
        tuple_key = tuple(cond_or_uncond) + tuple(x_like.shape)
        while control is not None:
            param_id += 1
            ph, pw = self.h * LATENT_SCALE_FACTOR, self.w * LATENT_SCALE_FACTOR

            if self.control_params.get(tuple_key) is None:
                self.control_params[tuple_key] = []
            val = self.control_params[tuple_key]
            while len(val) <= param_id:
                val.append([])
            while len(val[param_id]) <= batch_id:
                val[param_id].append(None)

            cached_hint = val[param_id][batch_id]
            hint_original = getattr(control, "cond_hint_original", None)
            if cached_hint is None and hint_original is not None:
                dtype = getattr(control, "manual_cast_dtype", None) or getattr(getattr(control, "control_model", None), "dtype", None) or x_like.dtype
                cond_hint = hint_original
                if cond_hint.shape[-2:] != (ph, pw):
                    cond_hint = _resize_hint(cond_hint, pw, ph, x_like)
                cond_hint = cond_hint.to(device=getattr(control, "device", x_like.device), dtype=dtype)
                if cond_hint.shape[0] < batch_size:
                    cond_hint = self.repeat_tensor(cond_hint, ceildiv(batch_size, cond_hint.shape[0]))[:batch_size]
                tiles = [cond_hint[:, :, bbox[1] * LATENT_SCALE_FACTOR : bbox[3] * LATENT_SCALE_FACTOR, bbox[0] * LATENT_SCALE_FACTOR : bbox[2] * LATENT_SCALE_FACTOR] for bbox in bboxes]
                cached_hint = torch.cat(tiles, dim=0)
                val[param_id][batch_id] = cached_hint

            if cached_hint is not None:
                control.cond_hint = cached_hint
            control = getattr(control, "previous_controlnet", None)


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:
            if x_in.shape[2] != 1:
                raise AssertionError("Expected a singleton frame dimension for 5D latents.")
            x_in = x_in.squeeze(2)

        n, _, h, w = x_in.shape
        if self.weights is None or self.h != h or self.w != w:
            self.h, self.w = h, w
            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:
                            if icond.shape[2] != 1:
                                raise AssertionError("Expected a singleton frame dimension for 5D conditioning.")
                            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, c_in, cond_or_uncond, bboxes, n, batch_id)
                    c_tile["control"] = c_in["control"].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 = getattr(control, "previous_controlnet", None)

                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, :, :, :]

        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):
    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:
            if x_in.shape[2] != 1:
                raise AssertionError("Expected a singleton frame dimension for 5D latents.")
            x_in = x_in.squeeze(2)

        n, _, h, w = x_in.shape
        if self.weights is None or self.h != h or self.w != w:
            self.h, self.w = h, w
            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 = []
                icond_map = {}
                for bbox in bboxes:
                    x_tile_list.append(x_in[bbox.slicer])
                    for key in ["y", "c_concat"]:
                        if key in c_in:
                            icond = c_in[key]
                            if icond.ndim == 5:
                                if icond.shape[2] != 1:
                                    raise AssertionError("Expected a singleton frame dimension for 5D conditioning.")
                                icond = icond.squeeze(2)

                            if icond.shape[2:] == (self.h, self.w):
                                icond = icond[bbox.slicer]
                            icond_map.setdefault(key, []).append(icond)

                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 ["y", "c_concat"]:
                    if key in icond_map:
                        c_tile[key] = torch.cat(icond_map[key], dim=0)
                if "control" in c_in:
                    self.process_controlnet(x_tile, c_in, cond_or_uncond, bboxes, n, batch_id)
                    c_tile["control"] = c_in["control"].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 = getattr(control, "previous_controlnet", None)

                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_tile = self.tile_weights * self.rescale_factor[bbox.slicer]
                    self.x_buffer[bbox.slicer] += x_tile_out[i * n : (i + 1) * n, :, :, :] * w_tile

        x_out = self.x_buffer
        if is_5d:
            x_out = x_out.unsqueeze(2)
        return x_out


class TiledDiffusion:
    @staticmethod
    def apply(model, method: str, tile_width: int, tile_height: int, tile_overlap: int, tile_batch_size: int):
        if method == "MultiDiffusion":
            implement = MultiDiffusion()
        elif method == "Mixture of Diffusers":
            implement = MixtureOfDiffusers()
        else:
            raise ValueError(f"Unsupported tiled diffusion method: {method}")

        implement.tile_width = max(1, tile_width // LATENT_SCALE_FACTOR)
        implement.tile_height = max(1, tile_height // LATENT_SCALE_FACTOR)
        implement.tile_overlap = max(0, tile_overlap // LATENT_SCALE_FACTOR)
        implement.tile_batch_size = max(1, tile_batch_size)

        model_clone = model.clone()
        model_clone.set_model_unet_function_wrapper(implement)
        model_clone.model_options["tiled_diffusion"] = {
            "method": method,
            "tile_width": tile_width,
            "tile_height": tile_height,
            "tile_overlap": tile_overlap,
            "tile_batch_size": tile_batch_size,
        }
        return model_clone