# adapted from https://github.com/kohya-ss/ControlNet-LLLite-ComfyUI
# basically, all the LLLite core code is from there, which I then combined with
# Advanced-ControlNet features and QoL
import math
from typing import Union
from torch import Tensor
import torch
import os

import comfy.utils
import comfy.ops
import comfy.model_management
from comfy.model_patcher import ModelPatcher
from comfy.controlnet import ControlBase

from .logger import logger
from .utils import (AdvancedControlBase, TimestepKeyframeGroup, ControlWeights, broadcast_image_to_extend, extend_to_batch_size,
                    deepcopy_with_sharing, prepare_mask_batch)


# based on set_model_patch code in comfy/model_patcher.py
def set_model_patch(model_options, patch, name):
    to = model_options["transformer_options"]
    # check if patch was already added
    if "patches" in to:
        current_patches = to["patches"].get(name, [])
        if patch in current_patches:
            return
    if "patches" not in to:
        to["patches"] = {}
    to["patches"][name] = to["patches"].get(name, []) + [patch]

def set_model_attn1_patch(model_options, patch):
    set_model_patch(model_options, patch, "attn1_patch")

def set_model_attn2_patch(model_options, patch):
    set_model_patch(model_options, patch, "attn2_patch")


def extra_options_to_module_prefix(extra_options):
    # extra_options = {'transformer_index': 2, 'block_index': 8, 'original_shape': [2, 4, 128, 128], 'block': ('input', 7), 'n_heads': 20, 'dim_head': 64}

    # block is: [('input', 4), ('input', 5), ('input', 7), ('input', 8), ('middle', 0),
    #   ('output', 0), ('output', 1), ('output', 2), ('output', 3), ('output', 4), ('output', 5)]
    # transformer_index is: [0, 1, 2, 3, 4, 5, 6, 7, 8], for each block
    # block_index is: 0-1 or 0-9, depends on the block
    # input 7 and 8, middle has 10 blocks

    # make module name from extra_options
    block = extra_options["block"]
    block_index = extra_options["block_index"]
    if block[0] == "input":
        module_pfx = f"lllite_unet_input_blocks_{block[1]}_1_transformer_blocks_{block_index}"
    elif block[0] == "middle":
        module_pfx = f"lllite_unet_middle_block_1_transformer_blocks_{block_index}"
    elif block[0] == "output":
        module_pfx = f"lllite_unet_output_blocks_{block[1]}_1_transformer_blocks_{block_index}"
    else:
        raise Exception(f"ControlLLLite: invalid block name '{block[0]}'. Expected 'input', 'middle', or 'output'.")
    return module_pfx


class LLLitePatch:
    ATTN1 = "attn1"
    ATTN2 = "attn2"
    def __init__(self, modules: dict[str, 'LLLiteModule'], patch_type: str, control: Union[AdvancedControlBase, ControlBase]=None):
        self.modules = modules
        self.control = control
        self.patch_type = patch_type
        #logger.error(f"create LLLitePatch: {id(self)},{control}")
    
    def __call__(self, q, k, v, extra_options):
        #logger.error(f"in __call__: {id(self)}")
        # determine if have anything to run
        if self.control.timestep_range is not None:
            # it turns out comparing single-value tensors to floats is extremely slow
            # a: Tensor = extra_options["sigmas"][0]
            if self.control.t > self.control.timestep_range[0] or self.control.t < self.control.timestep_range[1]:
                return q, k, v

        module_pfx = extra_options_to_module_prefix(extra_options)

        is_attn1 = q.shape[-1] == k.shape[-1]  # self attention
        if is_attn1:
            module_pfx = module_pfx + "_attn1"
        else:
            module_pfx = module_pfx + "_attn2"

        module_pfx_to_q = module_pfx + "_to_q"
        module_pfx_to_k = module_pfx + "_to_k"
        module_pfx_to_v = module_pfx + "_to_v"

        if module_pfx_to_q in self.modules:
            q = q + self.modules[module_pfx_to_q](q, self.control)
        if module_pfx_to_k in self.modules:
            k = k + self.modules[module_pfx_to_k](k, self.control)
        if module_pfx_to_v in self.modules:
            v = v + self.modules[module_pfx_to_v](v, self.control)

        return q, k, v

    def to(self, device):
        #logger.info(f"to... has control? {self.control}")
        for d in self.modules.keys():
            self.modules[d] = self.modules[d].to(device)
        return self
    
    def set_control(self, control: Union[AdvancedControlBase, ControlBase]) -> 'LLLitePatch':
        self.control = control
        return self
        #logger.error(f"set control for LLLitePatch: {id(self)}, cn: {id(control)}")

    def clone_with_control(self, control: AdvancedControlBase):
        #logger.error(f"clone-set control for LLLitePatch: {id(self)},{id(control)}")
        return LLLitePatch(self.modules, self.patch_type, control)

    def cleanup(self):
        #total_cleaned = 0
        for module in self.modules.values():
            module.cleanup()
        #    total_cleaned += 1
        #logger.info(f"cleaned modules: {total_cleaned}, {id(self)}")
        #logger.error(f"cleanup LLLitePatch: {id(self)}")

    # make sure deepcopy does not copy control, and deepcopied LLLitePatch should be assigned to control
    # def __deepcopy__(self, memo):
    #     self.cleanup()
    #     to_return: LLLitePatch = deepcopy_with_sharing(self, shared_attribute_names = ['control'], memo=memo)
    #     #logger.warn(f"patch {id(self)} turned into {id(to_return)}")
    #     try:
    #         if self.patch_type == self.ATTN1:
    #             to_return.control.patch_attn1 = to_return
    #         elif self.patch_type == self.ATTN2:
    #             to_return.control.patch_attn2 = to_return
    #     except Exception:
    #         pass
    #     return to_return


# TODO: use comfy.ops to support fp8 properly
class LLLiteModule(torch.nn.Module):
    def __init__(
        self,
        name: str,
        is_conv2d: bool,
        in_dim: int,
        depth: int,
        cond_emb_dim: int,
        mlp_dim: int,
    ):
        super().__init__()
        self.name = name
        self.is_conv2d = is_conv2d
        self.is_first = False

        modules = []
        modules.append(torch.nn.Conv2d(3, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))  # to latent (from VAE) size*2
        if depth == 1:
            modules.append(torch.nn.ReLU(inplace=True))
            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
        elif depth == 2:
            modules.append(torch.nn.ReLU(inplace=True))
            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=4, stride=4, padding=0))
        elif depth == 3:
            # kernel size 8 is too large, so set it to 4
            modules.append(torch.nn.ReLU(inplace=True))
            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))
            modules.append(torch.nn.ReLU(inplace=True))
            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))

        self.conditioning1 = torch.nn.Sequential(*modules)

        if self.is_conv2d:
            self.down = torch.nn.Sequential(
                torch.nn.Conv2d(in_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
                torch.nn.ReLU(inplace=True),
            )
            self.mid = torch.nn.Sequential(
                torch.nn.Conv2d(mlp_dim + cond_emb_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
                torch.nn.ReLU(inplace=True),
            )
            self.up = torch.nn.Sequential(
                torch.nn.Conv2d(mlp_dim, in_dim, kernel_size=1, stride=1, padding=0),
            )
        else:
            self.down = torch.nn.Sequential(
                torch.nn.Linear(in_dim, mlp_dim),
                torch.nn.ReLU(inplace=True),
            )
            self.mid = torch.nn.Sequential(
                torch.nn.Linear(mlp_dim + cond_emb_dim, mlp_dim),
                torch.nn.ReLU(inplace=True),
            )
            self.up = torch.nn.Sequential(
                torch.nn.Linear(mlp_dim, in_dim),
            )

        self.depth = depth
        self.cond_emb = None
        self.cx_shape = None
        self.prev_batch = 0
        self.prev_sub_idxs = None

    def cleanup(self):
        del self.cond_emb
        self.cond_emb = None
        self.cx_shape = None
        self.prev_batch = 0
        self.prev_sub_idxs = None

    def forward(self, x: Tensor, control: Union[AdvancedControlBase, ControlBase]):
        mask = None
        mask_tk = None
        #logger.info(x.shape)
        if self.cond_emb is None or control.sub_idxs != self.prev_sub_idxs or x.shape[0] != self.prev_batch:
            # print(f"cond_emb is None, {self.name}")
            cond_hint = control.cond_hint.to(x.device, dtype=x.dtype)
            if control.latent_dims_div2 is not None and x.shape[-1] != 1280:
                cond_hint = comfy.utils.common_upscale(cond_hint, control.latent_dims_div2[0] * 8, control.latent_dims_div2[1] * 8, 'nearest-exact', "center").to(x.device, dtype=x.dtype)
            elif control.latent_dims_div4 is not None and x.shape[-1] == 1280:
                cond_hint = comfy.utils.common_upscale(cond_hint, control.latent_dims_div4[0] * 8, control.latent_dims_div4[1] * 8, 'nearest-exact', "center").to(x.device, dtype=x.dtype)
            cx = self.conditioning1(cond_hint)
            self.cx_shape = cx.shape
            if not self.is_conv2d:
                # reshape / b,c,h,w -> b,h*w,c
                n, c, h, w = cx.shape
                cx = cx.view(n, c, h * w).permute(0, 2, 1)
            self.cond_emb = cx
        # save prev values
        self.prev_batch = x.shape[0]
        self.prev_sub_idxs = control.sub_idxs

        cx: torch.Tensor = self.cond_emb
        # print(f"forward {self.name}, {cx.shape}, {x.shape}")

        # TODO: make masks work for conv2d (could not find any ControlLLLites at this time that use them)
        # create masks
        if not self.is_conv2d:
            n, c, h, w = self.cx_shape
            if control.mask_cond_hint is not None:
                mask = prepare_mask_batch(control.mask_cond_hint, (1, 1, h, w)).to(cx.dtype)
                mask = mask.view(mask.shape[0], 1, h * w).permute(0, 2, 1)
            if control.tk_mask_cond_hint is not None:
                mask_tk = prepare_mask_batch(control.mask_cond_hint, (1, 1, h, w)).to(cx.dtype)
                mask_tk = mask_tk.view(mask_tk.shape[0], 1, h * w).permute(0, 2, 1)

        # x in uncond/cond doubles batch size
        if x.shape[0] != cx.shape[0]:
            if self.is_conv2d:
                cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1, 1)
            else:
                # print("x.shape[0] != cx.shape[0]", x.shape[0], cx.shape[0])
                cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1)
                if mask is not None:
                    mask = mask.repeat(x.shape[0] // mask.shape[0], 1, 1)
                if mask_tk is not None:
                    mask_tk = mask_tk.repeat(x.shape[0] // mask_tk.shape[0], 1, 1)

        if mask is None:
            mask = 1.0
        elif mask_tk is not None:
            mask = mask * mask_tk

        #logger.info(f"cs: {cx.shape}, x: {x.shape}, is_conv2d: {self.is_conv2d}")
        cx = torch.cat([cx, self.down(x)], dim=1 if self.is_conv2d else 2)
        cx = self.mid(cx)
        cx = self.up(cx)
        if control.latent_keyframes is not None:
            cx = cx * control.calc_latent_keyframe_mults(x=cx, batched_number=control.batched_number)
        if control.weights is not None and control.weights.has_uncond_multiplier:
            cond_or_uncond = control.batched_number.cond_or_uncond
            actual_length = cx.size(0) // control.batched_number
            for idx, cond_type in enumerate(cond_or_uncond):
                # if uncond, set to weight's uncond_multiplier
                if cond_type == 1:
                    cx[actual_length*idx:actual_length*(idx+1)] *= control.weights.uncond_multiplier
        return cx * mask * control.strength * control._current_timestep_keyframe.strength


class ControlLLLiteModules(torch.nn.Module):
    def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch):
        super().__init__()
        self.patch_attn1_modules = torch.nn.Sequential(*list(patch_attn1.modules.values()))
        self.patch_attn2_modules = torch.nn.Sequential(*list(patch_attn2.modules.values()))


class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase):
    # This ControlNet is more of an attention patch than a traditional controlnet
    def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch, timestep_keyframes: TimestepKeyframeGroup, device, ops: comfy.ops.disable_weight_init):
        super().__init__()
        AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite())
        self.device = device
        self.ops = ops
        self.patch_attn1 = patch_attn1.clone_with_control(self)
        self.patch_attn2 = patch_attn2.clone_with_control(self)
        self.control_model = ControlLLLiteModules(self.patch_attn1, self.patch_attn2)
        self.control_model_wrapped = ModelPatcher(self.control_model, load_device=device, offload_device=comfy.model_management.unet_offload_device())
        self.latent_dims_div2 = None
        self.latent_dims_div4 = None

    def live_model_patches(self, model_options):
        set_model_attn1_patch(model_options, self.patch_attn1.set_control(self))
        set_model_attn2_patch(model_options, self.patch_attn2.set_control(self))

    # def patch_model(self, model: ModelPatcher):
    #     model.set_model_attn1_patch(self.patch_attn1)
    #     model.set_model_attn2_patch(self.patch_attn2)

    def set_cond_hint_inject(self, *args, **kwargs):
        to_return = super().set_cond_hint_inject(*args, **kwargs)
        # cond hint for LLLite needs to be scaled between (-1, 1) instead of (0, 1)
        self.cond_hint_original = self.cond_hint_original * 2.0 - 1.0
        return to_return

    def pre_run_advanced(self, *args, **kwargs):
        AdvancedControlBase.pre_run_advanced(self, *args, **kwargs)
        #logger.error(f"in cn: {id(self.patch_attn1)},{id(self.patch_attn2)}")
        self.patch_attn1.set_control(self)
        self.patch_attn2.set_control(self)
        #logger.warn(f"in pre_run_advanced: {id(self)}")
    
    def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
        # normal ControlNet stuff
        control_prev = None
        if self.previous_controlnet is not None:
            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)

        if self.timestep_range is not None:
            if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
                return control_prev
        
        dtype = x_noisy.dtype
        # prepare cond_hint
        if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
            if self.cond_hint is not None:
                del self.cond_hint
            self.cond_hint = None
            # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
            if self.sub_idxs is not None:
                actual_cond_hint_orig = self.cond_hint_original
                if self.cond_hint_original.size(0) < self.full_latent_length:
                    actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length)
                self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(x_noisy.device)
            else:
                self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(x_noisy.device)
        if x_noisy.shape[0] != self.cond_hint.shape[0]:
            self.cond_hint = broadcast_image_to_extend(self.cond_hint, x_noisy.shape[0], batched_number)
        # some special logic here compared to other controlnets:
        # * The cond_emb in attn patches will divide latent dims by 2 or 4, integer
        # * Due to this loss, the cond_emb will become smaller than x input if latent dims are not divisble by 2 or 4
        divisible_by_2_h = x_noisy.shape[2]%2==0
        divisible_by_2_w = x_noisy.shape[3]%2==0
        if not (divisible_by_2_h and divisible_by_2_w):
            #logger.warn(f"{x_noisy.shape} not divisible by 2!")
            new_h = (x_noisy.shape[2]//2)*2
            new_w = (x_noisy.shape[3]//2)*2
            if not divisible_by_2_h:
                new_h += 2
            if not divisible_by_2_w:
                new_w += 2
            self.latent_dims_div2 = (new_h, new_w)
        divisible_by_4_h = x_noisy.shape[2]%4==0
        divisible_by_4_w =  x_noisy.shape[3]%4==0
        if not (divisible_by_4_h and divisible_by_4_w):
            #logger.warn(f"{x_noisy.shape} not divisible by 4!")
            new_h = (x_noisy.shape[2]//4)*4
            new_w = (x_noisy.shape[3]//4)*4
            if not divisible_by_4_h:
                new_h += 4
            if not divisible_by_4_w:
                new_w += 4
            self.latent_dims_div4 = (new_h, new_w)
        # prepare mask
        self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number)
        # done preparing; model patches will take care of everything now.
        # return normal controlnet stuff
        return control_prev
    
    def get_models(self):
        to_return: list = super().get_models()
        to_return.append(self.control_model_wrapped)
        return to_return

    def cleanup_advanced(self):
        super().cleanup_advanced()
        self.patch_attn1.cleanup()
        self.patch_attn2.cleanup()
        self.latent_dims_div2 = None
        self.latent_dims_div4 = None
    
    def copy(self):
        c = ControlLLLiteAdvanced(self.patch_attn1, self.patch_attn2, self.timestep_keyframes, self.device, self.ops)
        self.copy_to(c)
        self.copy_to_advanced(c)
        return c
    
    # deepcopy needs to properly keep track of objects to work between model.clone calls!
    # def __deepcopy__(self, *args, **kwargs):
    #     self.cleanup_advanced()
    #     return self

    # def get_models(self):
    #     # get_models is called once at the start of every KSampler run - use to reset already_patched status
    #     out = super().get_models()
    #     logger.error(f"in get_models! {id(self)}")
    #     return out


def load_controllllite(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, timestep_keyframe: TimestepKeyframeGroup=None):
    if controlnet_data is None:
        controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
    # adapted from https://github.com/kohya-ss/ControlNet-LLLite-ComfyUI
    # first, split weights for each module
    module_weights = {}
    for key, value in controlnet_data.items():
        fragments = key.split(".")
        module_name = fragments[0]
        weight_name = ".".join(fragments[1:])

        if module_name not in module_weights:
            module_weights[module_name] = {}
        module_weights[module_name][weight_name] = value

    unet_dtype = comfy.model_management.unet_dtype()
    load_device = comfy.model_management.get_torch_device()
    manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
    ops = comfy.ops.disable_weight_init
    if manual_cast_dtype is not None:
        ops = comfy.ops.manual_cast

    # next, load each module
    modules = {}
    for module_name, weights in module_weights.items():
        # kohya planned to do something about how these should be chosen, so I'm not touching this
        # since I am not familiar with the logic for this
        if "conditioning1.4.weight" in weights:
            depth = 3
        elif weights["conditioning1.2.weight"].shape[-1] == 4:
            depth = 2
        else:
            depth = 1

        module = LLLiteModule(
            name=module_name,
            is_conv2d=weights["down.0.weight"].ndim == 4,
            in_dim=weights["down.0.weight"].shape[1],
            depth=depth,
            cond_emb_dim=weights["conditioning1.0.weight"].shape[0] * 2,
            mlp_dim=weights["down.0.weight"].shape[0],
        )
        # load weights into module
        module.load_state_dict(weights)
        modules[module_name] = module.to(dtype=unet_dtype)
        if len(modules) == 1:
            module.is_first = True

    #logger.info(f"loaded {ckpt_path} successfully, {len(modules)} modules")

    patch_attn1 = LLLitePatch(modules=modules, patch_type=LLLitePatch.ATTN1)
    patch_attn2 = LLLitePatch(modules=modules, patch_type=LLLitePatch.ATTN2)
    control = ControlLLLiteAdvanced(patch_attn1=patch_attn1, patch_attn2=patch_attn2, timestep_keyframes=timestep_keyframe, device=load_device, ops=ops)
    return control