from __future__ import annotations

from dataclasses import dataclass
import re
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

import torch
import torch.nn as nn


PAPER_SPARSE_FLUX_COMBINED_BLOCKS: Tuple[int, ...] = (1, 15, 36, 41, 48)
PAPER_SOURCE_FLUX_DOUBLE_BLOCKS = 19
PAPER_SOURCE_FLUX_SINGLE_BLOCKS = 38
PAPER_SOURCE_FLUX_TOTAL_BLOCKS = PAPER_SOURCE_FLUX_DOUBLE_BLOCKS + PAPER_SOURCE_FLUX_SINGLE_BLOCKS
STATE_KEY = "diffaid_runtime_state"
_EPSILON = 1.0e-12


@dataclass(frozen=True)
class SharedConfig:
    strength: float
    sigma_start: float
    sigma_end: float
    sigma_ramp: float
    token_weight_mode: str
    token_tail: float
    cond_only: bool


@dataclass(frozen=True)
class FluxMappedBlocks:
    combined_1based: Tuple[int, ...]
    double_0based: Tuple[int, ...]
    single_0based: Tuple[int, ...]
    total_double: int
    total_single: int

    @property
    def total(self) -> int:
        return self.total_double + self.total_single


@dataclass(frozen=True)
class WanMappedBlocks:
    requested_1based: Tuple[int, ...]
    mapped_1based: Tuple[int, ...]
    mapped_0based: Tuple[int, ...]
    total: int


@dataclass(frozen=True)
class SdxlTargetSpec:
    stage: str
    block_number: int
    transformer_index: Optional[int] = None


def _dedupe_preserve_order(values: Iterable[int]) -> List[int]:
    seen = set()
    out: List[int] = []
    for value in values:
        if value in seen:
            continue
        seen.add(value)
        out.append(value)
    return out


def _parse_combined_block_indices(text: str) -> List[int]:
    parts = re.split(r"[\s,;:|]+", text.strip())
    out: List[int] = []
    for part in parts:
        if not part:
            continue
        try:
            value = int(part)
        except ValueError as exc:
            raise ValueError(f"Invalid block index '{part}'. Use a comma-separated list of positive integers.") from exc
        if value <= 0:
            raise ValueError(f"Invalid block index '{value}'. Indices are 1-based and must be >= 1.")
        out.append(value)
    if not out:
        raise ValueError("No block indices were provided.")
    return _dedupe_preserve_order(out)


def _parse_sdxl_targets(text: str) -> Tuple[SdxlTargetSpec, ...]:
    if not text.strip():
        return tuple()

    out: List[SdxlTargetSpec] = []
    for raw_part in re.split(r"[\n,;|]+", text.strip()):
        part = raw_part.strip()
        if not part:
            continue
        pieces = [p.strip().lower() for p in part.split(":") if p.strip()]
        if len(pieces) not in (2, 3):
            raise ValueError(f"Invalid SDXL target '{part}'. Use 'input:4', 'middle:0', or 'output:7:1'.")
        stage = pieces[0]
        if stage not in {"input", "middle", "output"}:
            raise ValueError(f"Invalid SDXL stage '{stage}' in target '{part}'.")
        try:
            block_number = int(pieces[1])
        except ValueError as exc:
            raise ValueError(f"Invalid block number in target '{part}'.") from exc
        if block_number < 0:
            raise ValueError(f"Invalid block number in target '{part}'. Must be >= 0.")
        transformer_index = None
        if len(pieces) == 3:
            try:
                transformer_index = int(pieces[2])
            except ValueError as exc:
                raise ValueError(f"Invalid transformer index in target '{part}'.") from exc
            if transformer_index < 0:
                raise ValueError(f"Invalid transformer index in target '{part}'. Must be >= 0.")
        out.append(SdxlTargetSpec(stage=stage, block_number=block_number, transformer_index=transformer_index))
    return tuple(out)


def _remap_stage_indices(indices_1based: Sequence[int], source_total: int, target_total: int) -> List[int]:
    if source_total <= 0:
        raise ValueError(f"source_total must be positive, got {source_total}.")
    if target_total <= 0:
        raise ValueError(f"target_total must be positive, got {target_total}.")

    out: List[int] = []
    for index_1based in indices_1based:
        if index_1based <= 0 or index_1based > source_total:
            raise ValueError(f"Stage-local block index {index_1based} is outside the source range 1..{source_total}.")
        if source_total == 1 or target_total == 1:
            out.append(1)
            continue
        scaled = round((index_1based - 1) * (target_total - 1) / (source_total - 1)) + 1
        out.append(int(scaled))
    return _dedupe_preserve_order(out)


def _paper_sparse_flux_source_double_indices() -> List[int]:
    return [index for index in PAPER_SPARSE_FLUX_COMBINED_BLOCKS if index <= PAPER_SOURCE_FLUX_DOUBLE_BLOCKS]


def _paper_sparse_flux_source_single_indices() -> List[int]:
    return [
        index - PAPER_SOURCE_FLUX_DOUBLE_BLOCKS
        for index in PAPER_SPARSE_FLUX_COMBINED_BLOCKS
        if index > PAPER_SOURCE_FLUX_DOUBLE_BLOCKS
    ]


def _remap_paper_sparse_flux_double_only_indices(total_double: int) -> List[int]:
    return _remap_stage_indices(
        _paper_sparse_flux_source_double_indices(),
        PAPER_SOURCE_FLUX_DOUBLE_BLOCKS,
        total_double,
    )


def _remap_paper_sparse_flux_indices(total_double: int, total_single: int) -> List[int]:
    mapped_double = _remap_paper_sparse_flux_double_only_indices(total_double)
    mapped_single_local = _remap_stage_indices(
        _paper_sparse_flux_source_single_indices(),
        PAPER_SOURCE_FLUX_SINGLE_BLOCKS,
        total_single,
    )
    mapped_single_combined = [total_double + index for index in mapped_single_local]
    return _dedupe_preserve_order([*mapped_double, *mapped_single_combined])


def _remap_paper_sparse_flux_to_single_block_list(total_blocks: int) -> List[int]:
    return _remap_stage_indices(
        PAPER_SPARSE_FLUX_COMBINED_BLOCKS,
        PAPER_SOURCE_FLUX_TOTAL_BLOCKS,
        total_blocks,
    )


def _smoothstep01(x: torch.Tensor) -> torch.Tensor:
    x = x.clamp(0.0, 1.0)
    return x * x * (3.0 - 2.0 * x)


def _sigma_window_gain(normalized_sigma: torch.Tensor, start: float, end: float, ramp: float) -> torch.Tensor:
    normalized_sigma = normalized_sigma.float().reshape(-1).clamp(0.0, 1.0)
    if start <= 0.0 and end >= 1.0:
        return torch.ones_like(normalized_sigma)
    if ramp <= 0.0:
        return ((normalized_sigma >= start) & (normalized_sigma <= end)).to(dtype=normalized_sigma.dtype)

    ramp = max(float(ramp), _EPSILON)
    left = _smoothstep01((normalized_sigma - (start - ramp)) / (2.0 * ramp))
    right = 1.0 - _smoothstep01((normalized_sigma - (end - ramp)) / (2.0 * ramp))
    return (left * right).clamp(0.0, 1.0)


def _token_weights(count: int, mode: str, tail: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    if count <= 0:
        raise ValueError("Token count must be positive.")
    if mode == "none" or count == 1:
        return torch.ones((count,), device=device, dtype=dtype)

    positions = torch.linspace(0.0, 1.0, steps=count, device=device, dtype=dtype)
    tail = max(0.0, min(1.0, float(tail)))
    if mode == "linear":
        return 1.0 - (1.0 - tail) * positions
    if mode == "exponential":
        safe_tail = max(tail, 1.0e-6)
        return torch.exp(torch.log(torch.tensor(safe_tail, device=device, dtype=dtype)) * positions)
    raise ValueError(f"Unsupported token weighting mode: {mode}")


def _coerce_batch_value(value, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    if torch.is_tensor(value):
        t = value.to(device=device, dtype=dtype).reshape(-1)
    else:
        t = torch.tensor([float(value)], device=device, dtype=dtype)
    if t.numel() == 1:
        return t.repeat(batch_size)
    if t.numel() == batch_size:
        return t
    return t[:1].repeat(batch_size)


def _as_float_tensor(value, device: torch.device) -> Optional[torch.Tensor]:
    if value is None:
        return None
    if torch.is_tensor(value):
        return value.detach().to(device=device, dtype=torch.float32).reshape(-1)
    try:
        return torch.tensor([float(value)], device=device, dtype=torch.float32)
    except (TypeError, ValueError):
        return None


def _has_tensor_payload(value) -> bool:
    if value is None:
        return False
    if torch.is_tensor(value):
        return value.numel() > 0
    if isinstance(value, (list, tuple, dict, str, bytes)):
        return len(value) > 0
    return True


def _detect_reference_latents(c: Dict, transformer_options: Dict) -> bool:
    for key in (
        "reference_latents",
        "ref_latents",
        "reference_latent",
        "ref_latent",
        "reference_image",
        "reference_images",
    ):
        if _has_tensor_payload(c.get(key)):
            return True
    ref_tokens = transformer_options.get("reference_image_num_tokens", 0)
    try:
        return int(ref_tokens) > 0
    except (TypeError, ValueError):
        return False


def _sequence_length(value) -> Optional[int]:
    if not torch.is_tensor(value):
        return None
    # WAN image context features are expected to be batched sequences
    # such as [batch, tokens, channels]. A 2D tensor is ambiguous because
    # shape[-2] would usually be the batch dimension, not a token count.
    if value.ndim < 3:
        return None
    try:
        return int(value.shape[-2])
    except (TypeError, ValueError):
        return None


def _detect_wan_context_image_tokens(c: Dict, transformer_options: Dict) -> int:
    for key in ("wan_context_img_len", "context_img_len", "context_image_token_count", "image_context_token_count"):
        value = transformer_options.get(key, None)
        try:
            count = int(value)
        except (TypeError, ValueError):
            continue
        if count > 0:
            return count

    for key in ("clip_fea", "clip_vision_output", "image_embeds"):
        count = _sequence_length(c.get(key))
        if count is not None and count > 0:
            return count
    return 0


def _cond_branch_gain(batch_size: int, config: SharedConfig, transformer_options: Optional[Dict], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    if not config.cond_only or transformer_options is None:
        return torch.ones((batch_size,), device=device, dtype=dtype)

    cond_or_uncond = transformer_options.get("cond_or_uncond", None)
    if cond_or_uncond is None:
        return torch.ones((batch_size,), device=device, dtype=dtype)

    if torch.is_tensor(cond_or_uncond):
        ids = cond_or_uncond.to(device=device).reshape(-1)
    else:
        try:
            ids = torch.tensor(list(cond_or_uncond), device=device)
        except TypeError:
            ids = torch.tensor([cond_or_uncond], device=device)
        ids = ids.reshape(-1)

    if ids.numel() == 0:
        return torch.ones((batch_size,), device=device, dtype=dtype)

    cond_rows = ids == 0
    if cond_rows.numel() == batch_size:
        mask = cond_rows
    elif batch_size % cond_rows.numel() == 0:
        mask = cond_rows.repeat_interleave(batch_size // cond_rows.numel())
    else:
        return torch.ones((batch_size,), device=device, dtype=dtype)
    return mask.to(dtype=dtype)


def _compute_alpha(reference_tensor: torch.Tensor, token_count: int, config: SharedConfig, transformer_options: Optional[Dict]) -> torch.Tensor:
    batch = reference_tensor.shape[0]
    transformer_options = transformer_options or {}
    normalized_sigma = transformer_options.get("normalized_sigma", None)
    if normalized_sigma is None:
        state = transformer_options.get(STATE_KEY, {}) or {}
        normalized_sigma = state.get("normalized_sigma", None)
    if normalized_sigma is None:
        normalized_sigma = 1.0

    sigma = _coerce_batch_value(normalized_sigma, batch, reference_tensor.device, reference_tensor.dtype)
    time_gain = _sigma_window_gain(sigma, config.sigma_start, config.sigma_end, config.sigma_ramp)
    token_gain = _token_weights(token_count, config.token_weight_mode, config.token_tail, reference_tensor.device, reference_tensor.dtype)
    branch_gain = _cond_branch_gain(batch, config, transformer_options, reference_tensor.device, reference_tensor.dtype)
    alpha = config.strength * time_gain[:, None, None] * branch_gain[:, None, None] * token_gain[None, :, None]
    return alpha.to(dtype=reference_tensor.dtype)


def _is_flux_family_model(model) -> bool:
    try:
        diffusion_model = model.get_model_object("diffusion_model")
    except Exception:
        return False
    required_attrs = ("double_blocks", "single_blocks", "txt_in", "forward_orig")
    return all(hasattr(diffusion_model, attr) for attr in required_attrs)


def _is_cross_attn_unet_model(model) -> bool:
    try:
        diffusion_model = model.get_model_object("diffusion_model")
    except Exception:
        return False
    return (
        hasattr(diffusion_model, "input_blocks")
        and hasattr(diffusion_model, "middle_block")
        and hasattr(diffusion_model, "output_blocks")
        and not hasattr(diffusion_model, "double_blocks")
    )


def _iter_wrapper_children(obj: Any):
    for attr in ("model", "diffusion_model", "inner_model", "module", "wrapped_model", "_orig_mod"):
        child = getattr(obj, attr, None)
        if child is not None and child is not obj:
            yield attr, child


def _looks_like_wan_model(inner: Any) -> bool:
    return (
        callable(getattr(inner, "forward_orig", None))
        and hasattr(inner, "blocks")
        and hasattr(inner, "patch_embedding")
        and hasattr(inner, "head")
        and not hasattr(inner, "double_blocks")
        and not hasattr(inner, "input_blocks")
    )


def _locate_wan_like_descendant(root: Any, root_name: str) -> Tuple[Optional[Any], Optional[str]]:
    if root is None:
        return None, root_name

    queue = [(root, root_name)]
    seen = set()
    fallback = None
    while queue:
        obj, name = queue.pop(0)
        obj_id = id(obj)
        if obj_id in seen:
            continue
        seen.add(obj_id)

        if _looks_like_wan_model(obj):
            return obj, name
        if fallback is None and callable(getattr(obj, "forward_orig", None)) and hasattr(obj, "blocks"):
            fallback = (obj, name)

        for attr, child in _iter_wrapper_children(obj):
            queue.append((child, f"{name}.{attr}"))

    if fallback is not None:
        return fallback
    return None, root_name


def _locate_wan_inner_model(model) -> Tuple[Optional[Any], Optional[str]]:
    try:
        diffusion_model = model.get_model_object("diffusion_model")
    except Exception:
        diffusion_model = None
    if diffusion_model is not None:
        inner, inner_name = _locate_wan_like_descendant(diffusion_model, "diffusion_model")
        if inner is not None:
            return inner, inner_name

    outer = getattr(model, "model", None)
    if outer is not None and hasattr(outer, "diffusion_model"):
        return _locate_wan_like_descendant(outer.diffusion_model, "model.diffusion_model")
    if hasattr(model, "diffusion_model"):
        return _locate_wan_like_descendant(model.diffusion_model, "diffusion_model")
    return None, None


def _is_wan_family_model(model) -> bool:
    inner, _inner_name = _locate_wan_inner_model(model)
    return inner is not None and hasattr(inner, "blocks")


def _get_flux_block_counts(model) -> Tuple[int, int]:
    diffusion_model = model.get_model_object("diffusion_model")
    return len(diffusion_model.double_blocks), len(diffusion_model.single_blocks)


def _get_wan_block_count(model) -> Tuple[int, Optional[str]]:
    inner, inner_name = _locate_wan_inner_model(model)
    if inner is None or not hasattr(inner, "blocks"):
        raise ValueError("This node expects a WAN-family MODEL whose diffusion model exposes a blocks list.")
    return len(inner.blocks), inner_name


def _map_indices_to_wan_blocks(requested_1based: Sequence[int], total_blocks: int) -> WanMappedBlocks:
    mapped_0based: List[int] = []
    for index_1based in requested_1based:
        if index_1based > total_blocks:
            raise ValueError(f"WAN block index {index_1based} is outside the model range 1..{total_blocks}.")
        mapped_0based.append(index_1based - 1)
    mapped_0based = _dedupe_preserve_order(mapped_0based)
    return WanMappedBlocks(
        requested_1based=tuple(_dedupe_preserve_order(requested_1based)),
        mapped_1based=tuple(index + 1 for index in mapped_0based),
        mapped_0based=tuple(mapped_0based),
        total=total_blocks,
    )


def _map_combined_indices_to_flux_stages(indices_1based: Sequence[int], total_double: int, total_single: int) -> FluxMappedBlocks:
    total = total_double + total_single
    double_indices: List[int] = []
    single_indices: List[int] = []
    for index_1based in indices_1based:
        if index_1based > total:
            raise ValueError(f"Block index {index_1based} is outside the model range 1..{total} (double={total_double}, single={total_single}).")
        if index_1based <= total_double:
            double_indices.append(index_1based - 1)
        else:
            single_indices.append(index_1based - total_double - 1)
    return FluxMappedBlocks(
        combined_1based=tuple(_dedupe_preserve_order(indices_1based)),
        double_0based=tuple(_dedupe_preserve_order(double_indices)),
        single_0based=tuple(_dedupe_preserve_order(single_indices)),
        total_double=total_double,
        total_single=total_single,
    )


class SharedTimestepWrapper(nn.Module):
    def __init__(self, existing_wrapper=None):
        super().__init__()
        self.existing_wrapper = existing_wrapper
        self._first_sigma_abs: Optional[float] = None
        self._last_sigma_abs: Optional[float] = None

    def _normalized_sigma(self, timestep, device: torch.device) -> torch.Tensor:
        t = _as_float_tensor(timestep, device=device)
        if t is None or t.numel() == 0:
            return torch.tensor([1.0], device=device, dtype=torch.float32)
        current = float(t.abs().max().item())
        reset_threshold = max(abs(self._last_sigma_abs or 0.0), 1.0) * 1.0e-4
        if self._first_sigma_abs is None or (self._last_sigma_abs is not None and current > self._last_sigma_abs + reset_threshold):
            self._first_sigma_abs = max(current, _EPSILON)
        self._last_sigma_abs = current
        if self._first_sigma_abs <= _EPSILON:
            return torch.zeros_like(t)
        return (t.abs() / self._first_sigma_abs).clamp(0.0, 1.0)

    def _inject_state(self, c: Dict, timestep) -> Dict:
        c = dict(c)
        transformer_options = dict(c.get("transformer_options", {}))
        state = dict(transformer_options.get(STATE_KEY, {}))
        device = timestep.device if torch.is_tensor(timestep) else torch.device("cpu")
        state["raw_timestep"] = timestep
        state["normalized_sigma"] = self._normalized_sigma(timestep, device=device)
        state["reference_latents"] = _detect_reference_latents(c, transformer_options)
        wan_context_img_len = _detect_wan_context_image_tokens(c, transformer_options)
        if wan_context_img_len > 0:
            state["wan_context_img_len"] = wan_context_img_len
        else:
            state.pop("wan_context_img_len", None)
        transformer_options[STATE_KEY] = state
        c["transformer_options"] = transformer_options
        return c

    def forward(self, model_function, params):
        def diffaid_model_function(input_x, timestep, **c):
            c = self._inject_state(c, timestep)
            return model_function(input_x, timestep, **c)

        if self.existing_wrapper is not None:
            return self.existing_wrapper(diffaid_model_function, params)
        return diffaid_model_function(params["input"], params["timestep"], **params["c"])


class FluxBlockReplacePatch(nn.Module):
    def __init__(self, stage: str, config: SharedConfig, existing_patch=None):
        super().__init__()
        self.stage = stage
        self.config = config
        self.existing_patch = existing_patch
        self._reported_reference_latents = False

    def _call_next(self, args: Dict, extra: Dict):
        if self.existing_patch is not None:
            return self.existing_patch(args, extra)
        return extra["original_block"](args)

    def _maybe_report_reference_latents(self, transformer_options: Dict):
        if self._reported_reference_latents:
            return
        state = transformer_options.get(STATE_KEY, {}) or {}
        if state.get("reference_latents", False):
            print("[ComfyUI-DiffAid-Patches] Reference latents detected; Diff-Aid modulation remains text-token-only.")
            self._reported_reference_latents = True

    def forward(self, args: Dict, extra: Dict):
        transformer_options = args.get("transformer_options", {}) or {}
        self._maybe_report_reference_latents(transformer_options)
        new_args = dict(args)

        if self.stage == "double":
            txt = args["txt"]
            alpha = _compute_alpha(txt, txt.shape[1], self.config, transformer_options)
            new_args["txt"] = txt + txt * alpha
            return self._call_next(new_args, extra)

        x = args["img"]
        img_slice = transformer_options.get("img_slice", None)
        if not img_slice:
            return self._call_next(new_args, extra)
        txt_len = int(img_slice[0])
        if txt_len <= 0:
            return self._call_next(new_args, extra)
        alpha = _compute_alpha(x[:, :txt_len, :], txt_len, self.config, transformer_options)
        prefix = x[:, :txt_len, :] + x[:, :txt_len, :] * alpha
        new_args["img"] = torch.cat((prefix, x[:, txt_len:, :]), dim=1)
        return self._call_next(new_args, extra)


class WanBlockReplacePatch(nn.Module):
    def __init__(self, config: SharedConfig, preserve_image_context_prefix: bool, existing_patch=None):
        super().__init__()
        self.config = config
        self.preserve_image_context_prefix = bool(preserve_image_context_prefix)
        self.existing_patch = existing_patch
        self._reported_reference_latents = False
        self._reported_image_context = False

    def _call_next(self, args: Dict, extra: Dict):
        # Diff-Aid mutates the WAN context first, then preserves any patch that already owned this block.
        if self.existing_patch is not None:
            return self.existing_patch(args, extra)
        return extra["original_block"](args)

    def _maybe_report_scope(self, transformer_options: Dict, text_start: int):
        state = transformer_options.get(STATE_KEY, {}) or {}
        if state.get("reference_latents", False) and not self._reported_reference_latents:
            print("[ComfyUI-DiffAid-Patches] WAN reference latents detected; Diff-Aid modulation remains context-token-only.")
            self._reported_reference_latents = True
        if text_start > 0 and not self._reported_image_context:
            print(f"[ComfyUI-DiffAid-Patches] WAN image-context prefix detected; preserving first {text_start} context tokens.")
            self._reported_image_context = True

    def _text_start(self, transformer_options: Dict, token_count: int) -> int:
        if not self.preserve_image_context_prefix:
            return 0
        state = transformer_options.get(STATE_KEY, {}) or {}
        try:
            text_start = int(state.get("wan_context_img_len", 0))
        except (TypeError, ValueError):
            text_start = 0
        return max(0, min(text_start, token_count))

    def forward(self, args: Dict, extra: Dict):
        transformer_options = args.get("transformer_options", {}) or {}
        key_name = "txt"
        context = args.get(key_name, None)
        if not torch.is_tensor(context):
            key_name = "context"
            context = args.get(key_name, None)
        if not torch.is_tensor(context) or context.ndim < 3 or context.shape[1] <= 0:
            return self._call_next(args, extra)

        token_count = int(context.shape[1])
        text_start = self._text_start(transformer_options, token_count)
        self._maybe_report_scope(transformer_options, text_start)
        if text_start >= token_count:
            return self._call_next(args, extra)

        text_tokens = context[:, text_start:, :]
        alpha = _compute_alpha(text_tokens, text_tokens.shape[1], self.config, transformer_options)
        text_mod = text_tokens + text_tokens * alpha

        new_args = dict(args)
        if text_start == 0:
            new_args[key_name] = text_mod
        else:
            new_args[key_name] = torch.cat((context[:, :text_start, :], text_mod), dim=1)
        return self._call_next(new_args, extra)


class SDXLCrossAttentionPatch(nn.Module):
    def __init__(self, config: SharedConfig, stage_filter: str, targets: Tuple[SdxlTargetSpec, ...]):
        super().__init__()
        self.config = config
        self.stage_filter = stage_filter
        self.targets = targets

    def _matches(self, extra_options: Dict) -> bool:
        block = extra_options.get("block", None)
        if not block or not isinstance(block, tuple) or len(block) < 2:
            return False
        stage = str(block[0]).lower()
        number = int(block[1])
        transformer_index = int(extra_options.get("block_index", 0))
        if self.targets:
            for spec in self.targets:
                if stage != spec.stage or number != spec.block_number:
                    continue
                if spec.transformer_index is not None and transformer_index != spec.transformer_index:
                    continue
                return True
            return False
        if self.stage_filter == "all":
            return True
        return stage == self.stage_filter

    def forward(self, n: torch.Tensor, context_attn2: torch.Tensor, value_attn2: Optional[torch.Tensor], extra_options: Dict):
        if context_attn2 is None or not self._matches(extra_options):
            return n, context_attn2, value_attn2
        alpha = _compute_alpha(context_attn2, context_attn2.shape[1], self.config, extra_options)
        context_mod = context_attn2 + context_attn2 * alpha
        if value_attn2 is None or value_attn2.shape == context_attn2.shape:
            value_source = context_attn2 if value_attn2 is None else value_attn2
            value_mod = value_source + value_source * alpha
        else:
            value_mod = value_attn2
        return n, context_mod, value_mod


def _get_existing_dit_replace_patch(model, block_kind: str, stage_index: int):
    # ComfyUI stores Flux and WAN block replacements under the same DiT namespace.
    transformer_options = model.model_options.get("transformer_options", {})
    patches_replace = transformer_options.get("patches_replace", {})
    dit_patches = patches_replace.get("dit", {})
    return dit_patches.get((block_kind, stage_index))


def _get_existing_flux_replace_patch(model, stage: str, stage_index: int):
    block_kind = "double_block" if stage == "double" else "single_block"
    return _get_existing_dit_replace_patch(model, block_kind, stage_index)


def _fmt_ints(values: Sequence[int]) -> str:
    return ", ".join(str(v) for v in values) if values else "-"


def _fmt_targets(values: Sequence[SdxlTargetSpec]) -> str:
    if not values:
        return "-"
    out = []
    for value in values:
        if value.transformer_index is None:
            out.append(f"{value.stage}:{value.block_number}")
        else:
            out.append(f"{value.stage}:{value.block_number}:{value.transformer_index}")
    return ", ".join(out)


class Flux2DiffAidSparsePatchNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "enabled": ("BOOLEAN", {"default": True}),
                "block_preset": (
                    ["paper_sparse_flux_double_only_safe", "paper_sparse_flux_full", "custom_combined_indices"],
                    {"default": "paper_sparse_flux_double_only_safe"},
                ),
                "block_indices": ("STRING", {"default": "1,15,36,41,48", "multiline": False, "advanced": True}),
                "strength": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step": 0.01}),
                "sigma_start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "sigma_end": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "sigma_ramp": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.5, "step": 0.001, "advanced": True}),
                "token_weight_mode": (["none", "linear", "exponential"], {"default": "none"}),
                "token_tail": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01, "advanced": True}),
                "apply_single_stream": ("BOOLEAN", {"default": False, "advanced": True}),
                "cond_only": ("BOOLEAN", {"default": True, "advanced": True}),
            }
        }

    RETURN_TYPES = ("MODEL", "STRING")
    RETURN_NAMES = ("model", "summary")
    FUNCTION = "patch"
    CATEGORY = "model_patches/diffaid"

    def patch(
        self,
        model,
        enabled: bool,
        block_preset: str,
        block_indices: str,
        strength: float,
        sigma_start: float,
        sigma_end: float,
        sigma_ramp: float,
        token_weight_mode: str,
        token_tail: float,
        apply_single_stream: bool = False,
        cond_only: bool = True,
    ):
        if not enabled:
            return model, "disabled"
        if not _is_flux_family_model(model):
            raise ValueError("This node only supports Flux-family MODEL objects that expose double_blocks/single_blocks. Use the SDXL node for cross-attention UNet models.")
        if sigma_start > sigma_end:
            raise ValueError(f"sigma_start ({sigma_start}) must be <= sigma_end ({sigma_end}).")

        total_double, total_single = _get_flux_block_counts(model)
        if block_preset == "paper_sparse_flux_double_only_safe":
            requested_combined_indices = list(PAPER_SPARSE_FLUX_COMBINED_BLOCKS)
            combined_indices = _remap_paper_sparse_flux_double_only_indices(total_double)
            preset_name = f"paper_sparse_flux_double_only_safe_remapped_from_flux1(source_double={PAPER_SOURCE_FLUX_DOUBLE_BLOCKS})"
        elif block_preset in {"paper_sparse_flux_full", "paper_sparse_flux"}:
            requested_combined_indices = list(PAPER_SPARSE_FLUX_COMBINED_BLOCKS)
            combined_indices = _remap_paper_sparse_flux_indices(total_double, total_single)
            legacy = "legacy_alias_" if block_preset == "paper_sparse_flux" else ""
            preset_name = f"{legacy}paper_sparse_flux_full_remapped_from_flux1(source_double={PAPER_SOURCE_FLUX_DOUBLE_BLOCKS},source_single={PAPER_SOURCE_FLUX_SINGLE_BLOCKS})"
        else:
            requested_combined_indices = _parse_combined_block_indices(block_indices)
            combined_indices = requested_combined_indices
            preset_name = "custom_combined_indices"

        mapped = _map_combined_indices_to_flux_stages(combined_indices, total_double, total_single)
        active_single_0based = mapped.single_0based if apply_single_stream else tuple()
        active_combined_1based = tuple([index + 1 for index in mapped.double_0based] + [mapped.total_double + index + 1 for index in active_single_0based])
        config = SharedConfig(float(strength), float(sigma_start), float(sigma_end), float(sigma_ramp), token_weight_mode, float(token_tail), bool(cond_only))

        patched = model.clone()
        existing_model_wrapper = patched.model_options.get("model_function_wrapper")
        patched.set_model_unet_function_wrapper(SharedTimestepWrapper(existing_wrapper=existing_model_wrapper))

        for block_index in mapped.double_0based:
            existing_patch = _get_existing_flux_replace_patch(patched, "double", block_index)
            patched.set_model_patch_replace(FluxBlockReplacePatch("double", config, existing_patch=existing_patch), "dit", "double_block", block_index)

        if apply_single_stream:
            for block_index in mapped.single_0based:
                existing_patch = _get_existing_flux_replace_patch(patched, "single", block_index)
                patched.set_model_patch_replace(FluxBlockReplacePatch("single", config, existing_patch=existing_patch), "dit", "single_block", block_index)

        inactive_single_0based = tuple(index for index in mapped.single_0based if index not in active_single_0based)
        summary = (
            f"flux_sparse preset={preset_name}; requested_combined_blocks=[{_fmt_ints(requested_combined_indices)}]; "
            f"mapped_combined_blocks=[{_fmt_ints(mapped.combined_1based)}]; active_combined_blocks=[{_fmt_ints(active_combined_1based)}]; "
            f"double_blocks_0based=[{_fmt_ints(mapped.double_0based)}]; single_blocks_0based=[{_fmt_ints(active_single_0based)}]; inactive_single_blocks_0based=[{_fmt_ints(inactive_single_0based)}]; "
            f"apply_single_stream={str(bool(apply_single_stream)).lower()}; cond_only={str(config.cond_only).lower()}; "
            f"strength={config.strength:.3f}; normalized_sigma_window=[{config.sigma_start:.3f}, {config.sigma_end:.3f}] ramp={config.sigma_ramp:.3f}; "
            f"token_weight_mode={config.token_weight_mode}; token_tail={config.token_tail:.3f}; reference_latents=runtime_detected_text_only; "
            f"model_total_blocks={mapped.total} (double={mapped.total_double}, single={mapped.total_single})"
        )
        print(f"[ComfyUI-DiffAid-Patches] {summary}")
        return patched, summary


class WanDiffAidSparsePatchNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "enabled": ("BOOLEAN", {"default": True}),
                "block_preset": (
                    ["paper_sparse_flux_remapped", "custom_block_indices"],
                    {"default": "paper_sparse_flux_remapped"},
                ),
                "block_indices": ("STRING", {"default": "1,15,36,41,48", "multiline": False, "advanced": True}),
                "strength": ("FLOAT", {"default": 0.35, "min": -1.0, "max": 1.0, "step": 0.01}),
                "sigma_start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "sigma_end": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "sigma_ramp": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.5, "step": 0.001, "advanced": True}),
                "token_weight_mode": (["none", "linear", "exponential"], {"default": "none"}),
                "token_tail": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01, "advanced": True}),
                "preserve_image_context_prefix": ("BOOLEAN", {"default": True, "advanced": True}),
                "cond_only": ("BOOLEAN", {"default": True, "advanced": True}),
            }
        }

    RETURN_TYPES = ("MODEL", "STRING")
    RETURN_NAMES = ("model", "summary")
    FUNCTION = "patch"
    CATEGORY = "model_patches/diffaid"

    def patch(
        self,
        model,
        enabled: bool,
        block_preset: str,
        block_indices: str,
        strength: float,
        sigma_start: float,
        sigma_end: float,
        sigma_ramp: float,
        token_weight_mode: str,
        token_tail: float,
        preserve_image_context_prefix: bool = True,
        cond_only: bool = True,
    ):
        if not enabled:
            return model, "disabled"
        if not _is_wan_family_model(model):
            raise ValueError("This node only supports WAN-family MODEL objects that expose a single blocks list. Use the Flux node for double_blocks/single_blocks models.")
        if sigma_start > sigma_end:
            raise ValueError(f"sigma_start ({sigma_start}) must be <= sigma_end ({sigma_end}).")

        total_blocks, inner_name = _get_wan_block_count(model)
        if block_preset == "paper_sparse_flux_remapped":
            requested_indices = list(PAPER_SPARSE_FLUX_COMBINED_BLOCKS)
            mapped_indices = _remap_paper_sparse_flux_to_single_block_list(total_blocks)
            preset_name = f"paper_sparse_flux_remapped_from_flux1(source_total={PAPER_SOURCE_FLUX_TOTAL_BLOCKS})"
        else:
            requested_indices = _parse_combined_block_indices(block_indices)
            mapped_indices = requested_indices
            preset_name = "custom_block_indices"

        mapped = _map_indices_to_wan_blocks(mapped_indices, total_blocks)
        config = SharedConfig(float(strength), float(sigma_start), float(sigma_end), float(sigma_ramp), token_weight_mode, float(token_tail), bool(cond_only))

        patched = model.clone()
        existing_model_wrapper = patched.model_options.get("model_function_wrapper")
        patched.set_model_unet_function_wrapper(SharedTimestepWrapper(existing_wrapper=existing_model_wrapper))

        for block_index in mapped.mapped_0based:
            existing_patch = _get_existing_dit_replace_patch(patched, "double_block", block_index)
            patched.set_model_patch_replace(
                WanBlockReplacePatch(config, preserve_image_context_prefix=preserve_image_context_prefix, existing_patch=existing_patch),
                "dit",
                "double_block",
                block_index,
            )

        summary = (
            f"wan_sparse preset={preset_name}; requested_blocks=[{_fmt_ints(requested_indices)}]; "
            f"mapped_blocks=[{_fmt_ints(mapped.mapped_1based)}]; blocks_0based=[{_fmt_ints(mapped.mapped_0based)}]; "
            f"cond_only={str(config.cond_only).lower()}; preserve_image_context_prefix={str(bool(preserve_image_context_prefix)).lower()}; "
            f"strength={config.strength:.3f}; normalized_sigma_window=[{config.sigma_start:.3f}, {config.sigma_end:.3f}] ramp={config.sigma_ramp:.3f}; "
            f"token_weight_mode={config.token_weight_mode}; token_tail={config.token_tail:.3f}; reference_latents=runtime_detected_context_only; "
            f"model_total_blocks={mapped.total}; inner={inner_name or '-'}"
        )
        print(f"[ComfyUI-DiffAid-Patches] {summary}")
        return patched, summary


class SDXLDiffAidCrossAttentionPatchNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "enabled": ("BOOLEAN", {"default": True}),
                "stage_filter": (["all", "input", "middle", "output"], {"default": "all"}),
                "block_targets": ("STRING", {"default": "", "multiline": False, "advanced": True}),
                "strength": ("FLOAT", {"default": 0.35, "min": -1.0, "max": 1.0, "step": 0.01}),
                "sigma_start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "sigma_end": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "sigma_ramp": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.5, "step": 0.001, "advanced": True}),
                "token_weight_mode": (["none", "linear", "exponential"], {"default": "linear"}),
                "token_tail": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01, "advanced": True}),
                "cond_only": ("BOOLEAN", {"default": True, "advanced": True}),
            }
        }

    RETURN_TYPES = ("MODEL", "STRING")
    RETURN_NAMES = ("model", "summary")
    FUNCTION = "patch"
    CATEGORY = "model_patches/diffaid"

    def patch(
        self,
        model,
        enabled: bool,
        stage_filter: str,
        block_targets: str,
        strength: float,
        sigma_start: float,
        sigma_end: float,
        sigma_ramp: float,
        token_weight_mode: str,
        token_tail: float,
        cond_only: bool = True,
    ):
        if not enabled:
            return model, "disabled"
        if not _is_cross_attn_unet_model(model):
            raise ValueError("This node expects an SDXL-style cross-attention UNet MODEL. It is not for Flux-family MMDiT models.")
        if sigma_start > sigma_end:
            raise ValueError(f"sigma_start ({sigma_start}) must be <= sigma_end ({sigma_end}).")

        targets = _parse_sdxl_targets(block_targets)
        config = SharedConfig(float(strength), float(sigma_start), float(sigma_end), float(sigma_ramp), token_weight_mode, float(token_tail), bool(cond_only))
        patched = model.clone()
        existing_model_wrapper = patched.model_options.get("model_function_wrapper")
        patched.set_model_unet_function_wrapper(SharedTimestepWrapper(existing_wrapper=existing_model_wrapper))
        patched.set_model_attn2_patch(SDXLCrossAttentionPatch(config=config, stage_filter=stage_filter, targets=targets))

        summary = (
            f"sdxl_cross_attention stage_filter={stage_filter}; block_targets=[{_fmt_targets(targets)}]; "
            f"strength={config.strength:.3f}; normalized_sigma_window=[{config.sigma_start:.3f}, {config.sigma_end:.3f}] ramp={config.sigma_ramp:.3f}; "
            f"token_weight_mode={config.token_weight_mode}; token_tail={config.token_tail:.3f}; cond_only={str(config.cond_only).lower()}"
        )
        print(f"[ComfyUI-DiffAid-Patches] {summary}")
        return patched, summary


NODE_CLASS_MAPPINGS = {
    "Flux2DiffAidSparsePatch": Flux2DiffAidSparsePatchNode,
    "WanDiffAidSparsePatch": WanDiffAidSparsePatchNode,
    "SDXLDiffAidCrossAttentionPatch": SDXLDiffAidCrossAttentionPatchNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Flux2DiffAidSparsePatch": "Flux-family Diff-Aid Sparse Patch",
    "WanDiffAidSparsePatch": "WAN Diff-Aid Sparse Patch",
    "SDXLDiffAidCrossAttentionPatch": "SDXL Diff-Aid Cross-Attention Patch",
}