import math
from dataclasses import dataclass
from typing import Any

import torch


ANIMA_COEFFICIENTS = (
    16.440769544691957,
    -21.14123745515793,
    13.006701290921894,
    -1.6731977662609427,
    0.08632681059623949,
)
_STEP_INFO_KEY = "anima_teacache_step_info"
_WRAPPER_KEY = "anima_teacache"


def poly1d(coefficients, value):
    result = 0.0
    for coefficient in coefficients:
        result = result * value + coefficient
    return result


@dataclass
class CacheEntry:
    previous_modulated_input: torch.Tensor | None = None
    previous_residual: torch.Tensor | None = None
    accumulated_distance: float = 0.0
    signature: tuple[Any, ...] | None = None


@dataclass
class CacheDecision:
    should_calc: bool
    cacheable: bool
    keys: list[Any]
    slices: list[slice]


class AnimaTeaCacheRuntime:
    def __init__(self, rel_l1_thresh, start_percent, end_percent, cache_device):
        self.rel_l1_thresh = rel_l1_thresh
        self.start_percent = start_percent
        self.end_percent = end_percent
        self.cache_device = cache_device
        self.entries: dict[Any, CacheEntry] = {}
        self.schedule_key = None
        self.last_step_index = None

    def reset(self):
        self.entries.clear()
        self.schedule_key = None
        self.last_step_index = None

    def diffusion_model_wrapper(self, executor, *args, **kwargs):
        diffusion_model = executor.class_obj
        transformer_options = kwargs.get("transformer_options", {})
        step_info = self._build_step_info(
            transformer_options.get("sigmas"),
            transformer_options.get("sample_sigmas"),
        )
        if step_info is not None:
            transformer_options[_STEP_INFO_KEY] = step_info
        executor.original = lambda *inner_args, **inner_kwargs: self.forward(
            diffusion_model, *inner_args, **inner_kwargs
        )
        return executor(*args, **kwargs)

    def _build_step_info(self, current_sigmas, sample_sigmas):
        if (
            not torch.is_tensor(current_sigmas)
            or not torch.is_tensor(sample_sigmas)
        ):
            return None

        sigmas = sample_sigmas.detach().flatten()
        total_steps = sigmas.numel() - 1
        if total_steps <= 0:
            return None

        current_sigma = current_sigmas.detach().flatten()[0].to(
            device=sigmas.device, dtype=sigmas.dtype
        )
        step_index = int((sigmas[:-1] - current_sigma).abs().argmin().item())
        percent = step_index / total_steps
        return {
            "index": step_index,
            "total": total_steps,
            "percent": percent,
            "schedule_key": tuple(sigmas.float().cpu().tolist()),
        }

    def _split_layout(self, tensor, transformer_options):
        labels = transformer_options.get("cond_or_uncond")
        uuids = transformer_options.get("uuids")
        if not labels:
            return [0], [slice(0, tensor.shape[0])]
        if tensor.shape[0] % len(labels) != 0:
            return None

        chunk_size = tensor.shape[0] // len(labels)
        occurrences = {}
        keys = []
        slices = []
        for index, label in enumerate(labels):
            uuid = uuids[index] if uuids and len(uuids) == len(labels) else None
            identity = (label, uuid)
            occurrence = occurrences.get(identity, 0)
            occurrences[identity] = occurrence + 1
            keys.append((label, uuid, occurrence))
            slices.append(slice(index * chunk_size, (index + 1) * chunk_size))
        return keys, slices

    def decide(self, modulated_input, transformer_options):
        step_info = transformer_options.get(_STEP_INFO_KEY)
        layout = self._split_layout(modulated_input, transformer_options)
        if step_info is None or layout is None:
            self.reset()
            return CacheDecision(True, False, [], [])

        step_index = step_info["index"]
        schedule_key = step_info["schedule_key"]
        if (
            self.schedule_key != schedule_key
            or (step_index == 0 and self.last_step_index != 0)
            or (self.last_step_index is not None and step_index < self.last_step_index)
        ):
            self.reset()
            self.schedule_key = schedule_key
        self.last_step_index = step_index

        keys, slices = layout
        in_range = self.start_percent <= step_info["percent"] <= self.end_percent
        boundary = step_index == 0 or step_index == step_info["total"] - 1
        force_full = not in_range or boundary

        current_inputs = []
        for key, section in zip(keys, slices):
            current = modulated_input[section].detach().to(self.cache_device)
            current_inputs.append(current)
            entry = self.entries.setdefault(key, CacheEntry())
            signature = (tuple(current.shape), current.dtype, tuple(modulated_input[section].stride()))

            if entry.signature != signature:
                entry.previous_modulated_input = None
                entry.previous_residual = None
                entry.accumulated_distance = 0.0
                entry.signature = signature
                force_full = True

            previous = entry.previous_modulated_input
            if previous is None or entry.previous_residual is None:
                force_full = True
            elif not force_full:
                denominator = previous.abs().mean()
                if not torch.isfinite(denominator) or denominator.item() <= 0.0:
                    force_full = True
                else:
                    relative_l1 = ((current - previous).abs().mean() / denominator).float().item()
                    scaled = poly1d(ANIMA_COEFFICIENTS, relative_l1)
                    if not math.isfinite(scaled):
                        force_full = True
                    else:
                        entry.accumulated_distance += scaled
                        if entry.accumulated_distance >= self.rel_l1_thresh:
                            force_full = True

        for key, current in zip(keys, current_inputs):
            self.entries[key].previous_modulated_input = current

        if force_full:
            for key in keys:
                self.entries[key].accumulated_distance = 0.0

        return CacheDecision(force_full, True, keys, slices)

    def commit_residual(self, decision, residual):
        if not decision.cacheable:
            return
        for key, section in zip(decision.keys, decision.slices):
            self.entries[key].previous_residual = residual[section].detach().to(self.cache_device)

    def reuse_residual(self, decision, hidden_states):
        residuals = []
        for key in decision.keys:
            residual = self.entries[key].previous_residual
            if residual is None:
                raise RuntimeError("Anima TeaCache residual is not initialized")
            residuals.append(residual.to(device=hidden_states.device, dtype=hidden_states.dtype))
        return hidden_states + torch.cat(residuals, dim=0)

    def _first_block_modulated_input(
        self,
        diffusion_model,
        hidden_states,
        timestep_embedding,
        adaln_lora,
        extra_pos_emb,
    ):
        block = diffusion_model.blocks[0]
        block_input = hidden_states
        if extra_pos_emb is not None:
            block_input = block_input + extra_pos_emb

        modulation = block.adaln_modulation_self_attn(timestep_embedding)
        if block.use_adaln_lora:
            modulation = modulation + adaln_lora
        shift, scale, _ = modulation.chunk(3, dim=-1)
        shift = shift[:, :, None, None, :]
        scale = scale[:, :, None, None, :]
        return block.layer_norm_self_attn(block_input) * (1 + scale) + shift

    def forward(
        self,
        diffusion_model,
        x,
        timesteps,
        context,
        fps=None,
        padding_mask=None,
        **kwargs,
    ):
        import comfy.ldm.common_dit

        original_shape = list(x.shape)
        x = comfy.ldm.common_dit.pad_to_patch_size(
            x,
            (
                diffusion_model.patch_temporal,
                diffusion_model.patch_spatial,
                diffusion_model.patch_spatial,
            ),
        )
        hidden_states, rope_embedding, extra_pos_emb = diffusion_model.prepare_embedded_sequence(
            x, fps=fps, padding_mask=padding_mask
        )

        if timesteps.ndim == 1:
            timesteps = timesteps.unsqueeze(1)
        sinusoidal_timestep = diffusion_model.t_embedder[0](timesteps).to(hidden_states.dtype)
        timestep_embedding, adaln_lora = diffusion_model.t_embedder[1](sinusoidal_timestep)
        timestep_embedding = diffusion_model.t_embedding_norm(timestep_embedding)

        diffusion_model.affline_scale_log_info = {
            "t_embedding_B_T_D": timestep_embedding.detach()
        }
        diffusion_model.affline_emb = timestep_embedding
        diffusion_model.crossattn_emb = context

        if extra_pos_emb is not None:
            assert hidden_states.shape == extra_pos_emb.shape

        # Keep ComfyUI model patches active when TeaCache takes over the
        # diffusion-model forward.  In particular, Anima LLLite prepares its
        # conditioning in a post-input patch and applies the conditioning to
        # attention/MLP inputs through the transformer patch collections.
        transformer_options = kwargs.get("transformer_options", {})
        patches = transformer_options.get("patches", {})
        if "post_input" in patches:
            transformer_options = transformer_options.copy()
            transformer_options["model_patch_data"] = {}
            for patch in patches["post_input"]:
                out = patch({
                    "img": hidden_states,
                    "x": x,
                    "transformer_options": transformer_options,
                })
                hidden_states = out["img"]

        block_kwargs = {
            "rope_emb_L_1_1_D": rope_embedding.unsqueeze(1).unsqueeze(0),
            "adaln_lora_B_T_3D": adaln_lora,
            "extra_per_block_pos_emb": extra_pos_emb,
            "transformer_options": transformer_options,
        }
        if hidden_states.dtype == torch.float16:
            hidden_states = hidden_states.float()

        modulated_input = self._first_block_modulated_input(
            diffusion_model,
            hidden_states,
            timestep_embedding,
            adaln_lora,
            extra_pos_emb,
        )
        transformer_options = block_kwargs["transformer_options"]
        decision = self.decide(modulated_input, transformer_options)

        if decision.should_calc:
            original_hidden_states = hidden_states
            for block_index, block in enumerate(diffusion_model.blocks):
                transformer_options["block_index"] = block_index
                hidden_states = block(
                    hidden_states,
                    timestep_embedding,
                    context,
                    **block_kwargs,
                )
            self.commit_residual(decision, hidden_states - original_hidden_states)
        else:
            hidden_states = self.reuse_residual(decision, hidden_states)

        output = diffusion_model.final_layer(
            hidden_states.to(context.dtype),
            timestep_embedding,
            adaln_lora_B_T_3D=adaln_lora,
        )
        output = diffusion_model.unpatchify(output)
        return output[
            :,
            :,
            : original_shape[-3],
            : original_shape[-2],
            : original_shape[-1],
        ]


class AnimaTeaCache:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL", {"tooltip": "The native ComfyUI Anima model."}),
                "rel_l1_thresh": (
                    "FLOAT",
                    {"default": 0.05, "min": 0.0, "max": 10.0, "step": 0.001},
                ),
                "start_percent": (
                    "FLOAT",
                    {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01},
                ),
                "end_percent": (
                    "FLOAT",
                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01},
                ),
                "cache_device": (["cuda", "cpu"], {"default": "cuda"}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "apply_teacache"
    CATEGORY = "TeaCache"
    TITLE = "Anima TeaCache"

    def apply_teacache(
        self,
        model,
        rel_l1_thresh,
        start_percent,
        end_percent,
        cache_device,
    ):
        if rel_l1_thresh == 0:
            return (model,)
        if start_percent > end_percent:
            raise ValueError("start_percent must not be greater than end_percent")

        from comfy.ldm.anima.model import Anima

        diffusion_model = model.get_model_object("diffusion_model")
        if not isinstance(diffusion_model, Anima):
            raise TypeError("Anima TeaCache only supports ComfyUI's native Anima model")

        import comfy.model_management as model_management
        import comfy.patcher_extension as patcher_extension

        WrappersMP = patcher_extension.WrappersMP

        device = (
            model_management.get_torch_device()
            if cache_device == "cuda"
            else torch.device("cpu")
        )
        runtime = AnimaTeaCacheRuntime(
            rel_l1_thresh=rel_l1_thresh,
            start_percent=start_percent,
            end_percent=end_percent,
            cache_device=device,
        )
        patched_model = model.clone()
        transformer_options = patched_model.model_options.setdefault("transformer_options", {})
        wrappers = transformer_options.setdefault("wrappers", {})
        wrappers.setdefault(WrappersMP.DIFFUSION_MODEL, {}).pop(_WRAPPER_KEY, None)
        patcher_extension.add_wrapper_with_key(
            WrappersMP.DIFFUSION_MODEL,
            _WRAPPER_KEY,
            runtime.diffusion_model_wrapper,
            patched_model.model_options,
            is_model_options=True,
        )
        return (patched_model,)


NODE_CLASS_MAPPINGS = {"AnimaTeaCache": AnimaTeaCache}
NODE_DISPLAY_NAME_MAPPINGS = {"AnimaTeaCache": AnimaTeaCache.TITLE}