import gc
import torch
import torch.nn.functional as F
import numpy as np
from torch import Tensor
from comfy.ldm.modules import attention as comfy_attention
from comfy.ldm.flux import math as flux_math
from comfy.ldm.flux import layers as flux_layers
from comfy.ldm.chroma import layers as chroma_layers
from comfy import model_management
from typing import List, Dict, Optional
from functools import partial
from einops import rearrange
from nodes import CLIPTextEncode
import latent_preview
import copy
orig_attention = flux_math.attention
orig_prepare_callback = latent_preview.prepare_callback

def convert_masks_to_patches(masks: torch.Tensor, pad) -> tuple:
    """
    Convert pixel masks to patch-level masks for attention.
    
    Args:
        masks: [N, H, W] tensor - subject regions in images
        
    Returns:
        patch_masks: [N, h*w] tensor where h=H_latent*8/16, w=W_latent*8/16
        h, w: patch dimensions
    """
    N = masks.shape[0]
    H, W = masks.shape[1], masks.shape[2]
    
    # Calculate patch dimensions (assuming 16x16 pixels per patch, but 1x1 in latent space)
    h = H // 16
    w = W // 16
    
    patch_masks = torch.zeros(N + pad, h, w, device=masks.device, dtype=masks.dtype)
    
    for i in range(N):
        # Convert to patch level using nearest-neighbor interpolation
        mask_resized = F.interpolate(
            masks[i:i+1, None, :, :], 
            size=(h, w), 
            mode='nearest-exact'
        )
        
        # Flatten to patch sequence
        patch_masks[i] = mask_resized[0][0]
    for i in range(pad):
        patch_masks[i + N] = torch.ones((h, w), device=masks.device, dtype=masks.dtype)
    return patch_masks, h, w

def match_condition(reference_conditions: List, target_condition: List) -> torch.Tensor:
    """
    Match reference conditions to target condition using dot product similarity.
    
    Args:
        reference_conditions: List of N conditioning tensors (one per reference subject)
        target_condition: Conditioning tensor for target images (shared)
        
    Returns:
        conditions_mask: 2D tensor [num_references, max_embedding_seq_length] 
                     with 1s where reference embeddings match target subsequences
    """
    # Extract embeddings
    ref_embeds = [cond[0][0] for cond in reference_conditions]  # text embeddings
    target_embed = target_condition[0][0]  # [1, seq_len, embed_dim]
    
    num_references = len(ref_embeds)
    target_seq_len = target_embed.shape[1]
    
    # Initialize mask tensor
    conditions_mask = torch.zeros(num_references, target_seq_len, dtype=target_embed.dtype, device=target_embed.device)
    
    # Process each reference embedding
    for i, ref_embed in enumerate(ref_embeds):
        ref_seq_len = ref_embed.shape[1]
        
        # If reference is longer than target, skip (shouldn't happen per spec)
        if ref_seq_len > target_seq_len:
            continue
            
        # Get reference embedding (remove batch dimension)
        ref_vec = ref_embed[0]  # [ref_seq_len, embed_dim]
        
        # Use unfold to get all subsequences of target with length ref_seq_len
        # target_embed[0] is [target_seq_len, embed_dim]
        # unfolded will be [target_seq_len - ref_seq_len + 1, ref_seq_len, embed_dim]
        unfolded = target_embed[0].unfold(0, ref_seq_len, 1).permute(0,2,1)  # [num_subseq, ref_seq_len, embed_dim]
        
        # Normalize for cosine similarity
        ref_normalized = F.normalize(ref_vec, p=2, dim=-1)  # [ref_seq_len, embed_dim]
        target_normalized = F.normalize(unfolded, p=2, dim=-1)  # [num_subseq, ref_seq_len, embed_dim]
        
        # Compute cosine similarity for all subsequences at once
        # Result: [num_subseq]
        similarities = torch.mean(torch.sum(ref_normalized.unsqueeze(0) * target_normalized, dim=-1), dim=-1)
        
        # Find the best matching subsequence
        best_match_idx = torch.argmax(similarities)
        
        # Mask the best matching subsequence
        start_idx = best_match_idx.item()
        conditions_mask[i, start_idx:start_idx + ref_seq_len] = 1.0
    
    return conditions_mask

def prepare_batched_conditions(reference_conditions: List, target_condition: List, num_targets) -> tuple:
    """
    Prepare the batched conditions for reference and target images.
    
    Args:
        reference_conditions: List of N conditioning tensors (one per reference subject)
        target_condition: Conditioning tensor for target images (shared)
        
    Returns:
        extended_condition: Batched conditioning tensor [B, num_reference + 1, max_seq, 4096]
        conditions_mask: Mask tensor [B, num_references + 1, max_seq]
        pooled_output: CLIP pooled output (None for Chroma, tensor for Flux)
    """
    # Extract text embeddings from conditions
    ref_embeds = []
    ref_pooled = []
    seq_lens = []
    for cond in reference_conditions:
        if cond is None:
            continue
        ref_embeds.append(cond[0][0])  # text embedding
        ref_pooled.append(cond[0][1].get('pooled_output', None))
    
    # Target condition
    target_embed = target_condition[0][0]
    target_pooled = target_condition[0][1].get('pooled_output', None)
    
    # Calculate max sequence length
    all_embeds = [target_embed] + ref_embeds
    max_seq_len = max(embed.shape[1] for embed in all_embeds)
    
    # Get batch size and embedding dimension
    embed_dim = target_embed.shape[2]  # Should be 4096
    
    # Create padded tensors with shape [B, num_reference + 1, max_seq, 4096]
    num_references = len(ref_embeds)
    padded_embeds = torch.zeros(num_references + num_targets, max_seq_len, embed_dim, dtype=target_embed.dtype, device=target_embed.device)
    conditions_mask = torch.zeros(num_references + num_targets, max_seq_len, dtype=target_embed.dtype, device=target_embed.device)
    
    # Fill in target embeds (first position)
    target_seq_len = target_embed.shape[1]

    
    # Fill in reference embeds
    for i, ref_embed in enumerate(ref_embeds):
        ref_seq_len = ref_embed.shape[1]
        padded_embeds[i, :ref_seq_len, :] = ref_embed[0]
        conditions_mask[i, :ref_seq_len] = 1.0
        seq_lens.append(ref_seq_len)
    
    for i in range(num_targets):
        padded_embeds[i + num_references, :target_seq_len, :] = target_embed[0]
        conditions_mask[i + num_references, :target_seq_len] = 1.0        
        seq_lens.append(target_seq_len)
    # Handle pooled output (only used by Flux)
    if target_pooled is not None:
        # For Flux, use the target pooled output
        pooled_output = target_pooled
    else:
        pooled_output = None
    
    extended_condition = [[
        padded_embeds,
        {'pooled_output': pooled_output}
    ]]
    
    return extended_condition, conditions_mask, max_seq_len, seq_lens


class ReferenceConditionCombine:
    RETURN_TYPES = ("CONDITIONING", "MASK")
    RETURN_NAMES = ("condition", "condition_mask")
    FUNCTION = "go"
    
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "num_targets": ("INT",),
            },
            "optional": {
                "reference_condition1": ("CONDITIONING",),
                "reference_condition2": ("CONDITIONING",),
                "reference_condition3": ("CONDITIONING",),
                "reference_condition4": ("CONDITIONING",),
                "reference_condition5": ("CONDITIONING",),
                "target_condition": ("CONDITIONING",),
            }
        }
    @classmethod
    def go(cls, *,
           num_targets, reference_condition1=None, reference_condition2=None,
           reference_condition3=None, reference_condition4=None, reference_condition5=None,
           target_condition=None):
        
        reference_conditions = [reference_condition1, reference_condition2, reference_condition3, 
                        reference_condition4, reference_condition5]
        extended_condition, conditions_mask, max_seq_len, seq_lens = prepare_batched_conditions(reference_conditions, target_condition, num_targets)
        return extended_condition, conditions_mask.unsqueeze(0)
    
class BatchedMaskedReferenceGen:
    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "go"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "source_patch_masks": ("MASK",),
                "target_patch_masks": ("MASK",),
                "conditions_mask": ("MASK",),
                "conditions_ref_mask":  ("MASK",),
            },
        }
    @classmethod
    def go(cls, *, model: object, source_patch_masks: torch.Tensor, target_patch_masks: torch.Tensor,
           conditions_mask: torch.Tensor, conditions_ref_mask=torch.Tensor):
        
        print(f'BatchedMaskedReferenceGen - Processing with model: {model}')
        if conditions_mask.dim() == 3:
            conditions_mask = conditions_mask[0]
        if conditions_ref_mask.dim() == 3:
            conditions_ref_mask = conditions_ref_mask[0]
        # Detect model type
        model_type = "flux"
        if model.model.__class__.__name__ == "Chroma":
            model_type = "chroma"
        print(f'Detected model type: {model_type}')

        num_references = conditions_ref_mask.shape[0]
        num_targets = conditions_mask.shape[0] - num_references

        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
        w = source_patch_masks.shape[-1]
        h = source_patch_masks.shape[1]
        patch_count = h * w
        source_patch_masks = source_patch_masks.flatten(start_dim=1)
        target_patch_masks = target_patch_masks.flatten(start_dim=1)
        reference_seq_mask = torch.log(torch.cat([conditions_mask[:num_references], torch.ones((num_references, patch_count), dtype=attn_dtype)], dim=1))
        source_seq_masks = torch.log(torch.cat([conditions_mask[:num_references+1], source_patch_masks], dim=1))
        target_seq_masks = torch.log(torch.cat([conditions_mask[num_references].unsqueeze(0).repeat(num_references+1,1), target_patch_masks], dim=1))

        def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor,  reference_seq_mask, target_seq_masks, source_seq_masks) -> Tensor:
            """
            Implementation using PyTorch's scaled_dot_product_attention with custom masking.
            
            Args:
                q: Query tensor of shape [B, H, L, D]
                k: Key tensor of shape [B, H, L, D]
                v: Value tensor of shape [B, H, L, D]
                pe: Positional encoding tensor
                attn_mask: Optional attention mask of shape [B, H, L, L]
                
            Returns:
                Output tensor of shape [B, L, H*D]
            """
            target_count = q.shape[0] - num_references
            # Apply rotary positional encoding for Flux models
            if pe is not None:
                q, k = flux_math.apply_rope(q, k, pe)
            
            # Get actual sequence length from QKV tensors
            actual_seq_len = q.shape[2]
            
            #if attn_mask is None or attn_mask.shape[-2] != actual_seq_len:
            #    this_source = torch.cat([torch.zeros([num_references, text_len]), source_patch_masks], dim=-1)
            #    this_target = torch.cat([torch.zeros([num_references, text_len]), source_patch_masks], dim=-1)
            #    attn_mask = torch.stack(torch.min(this_source.unsqueeze(1), this_target.unsqueeze(2)), dim=-1).to(device)
                
            reference_seq_mask_ = reference_seq_mask.to(device, dtype=attn_dtype)
            reference_attn_mask = torch.min(reference_seq_mask_.unsqueeze(1), reference_seq_mask_.unsqueeze(2))
            target_attn_mask = rearrange(torch.min(source_seq_masks.to(device, dtype=attn_dtype).unsqueeze(1), target_seq_masks.to(device, dtype=attn_dtype).unsqueeze(2)), "B H W -> 1 W (B H)")

            h = torch.zeros(q.shape, device=device)
                # PyTorch's scaled_dot_product_attention
            h[:num_references] = F.scaled_dot_product_attention(
                q[:num_references], k[:num_references], v[:num_references],
                attn_mask=reference_attn_mask
            )
                
            B, H, L, D = k.shape
            k = k.permute(1, 0, 2, 3).reshape(1, H, B * L, D)
            B, H, L, D = v.shape
            v = v.permute(1, 0, 2, 3).reshape(1, H, B * L, D)

            for i in range(target_count):
                target_idx = i + num_references
                if i > 0:
                    k[:, :, num_references*actual_seq_len : num_references*actual_seq_len + actual_seq_len] = k[:, :, target_idx*actual_seq_len : target_idx*actual_seq_len + actual_seq_len]
                    v[:, :, num_references*actual_seq_len : num_references*actual_seq_len + actual_seq_len] = v[:, :, target_idx*actual_seq_len : target_idx*actual_seq_len + actual_seq_len]
                
                h[target_idx:target_idx + 1] = F.scaled_dot_product_attention(
                    q[target_idx: target_idx + 1], 
                    k[:, :, :num_references*actual_seq_len + actual_seq_len], 
                    v[:, :, :num_references*actual_seq_len + actual_seq_len], attn_mask=target_attn_mask)
            
            # Reshape back to the expected output format
            h = rearrange(h, "B H L D -> B L (H D)")
            return h

        # Create the attention override function
        def batched_masked_attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask, transformer_options,  reference_seq_mask, target_seq_masks, source_seq_masks) -> Tensor:
            if q.shape[0] != num_references + num_targets or reference_seq_mask is None or reference_seq_mask.shape[-1] != q.shape[2]:
                return orig_attention(q, k, v, pe)
            else:
                return masked_attention(q, k, v, pe,  reference_seq_mask, target_seq_masks, source_seq_masks)
        
        # Apply the attention override based on model type
        overrider = partial(batched_masked_attention_override, reference_seq_mask=reference_seq_mask, target_seq_masks=target_seq_masks, source_seq_masks=source_seq_masks)
        if model_type == "flux":
            flux_math.attention = overrider
        else:  # chroma
            chroma_layers.attention = overrider
        
        # Clean up
        #del source_patch_masks, target_patch_masks
        gc.collect()
        torch.cuda.empty_cache()
        
        print('BatchedMaskedReferenceGen processing completed')
        
        return (model,)


class ConditioningMatchMask:
    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("mask",)
    FUNCTION = "match"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
            },
            "optional": {
                "reference_condition1": ("CONDITIONING",),
                "reference_condition2": ("CONDITIONING",),
                "reference_condition3": ("CONDITIONING",),
                "reference_condition4": ("CONDITIONING",),
                "reference_condition5": ("CONDITIONING",),
                "target_condition": ("CONDITIONING",),
            }
        }

    def match(self, reference_condition1=None, reference_condition2=None, reference_condition3=None, 
              reference_condition4=None, reference_condition5=None, target_condition=None):
        # Collect reference conditions
        reference_conditions = []
        ref_conditions = [reference_condition1, reference_condition2, reference_condition3, 
                         reference_condition4, reference_condition5]
        
        for cond in ref_conditions:
            if cond is not None:
                reference_conditions.append(cond)
        
        if not reference_conditions:
            raise ValueError("At least one reference condition must be provided")
        
        if target_condition is None:
            raise ValueError("Target condition must be provided")
        
        # Call the match_condition function
        conditions_mask = match_condition(reference_conditions, target_condition)
        
        return (conditions_mask.unsqueeze(0),)


class MasksToPatches:
    RETURN_TYPES = ("MASK", "INT", "INT")
    RETURN_NAMES = ("patch_masks", "h", "w")
    FUNCTION = "convert"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "masks": ("MASK",),
                "target_count": ("INT",{"default":1,"min":0}),
            }
        }

    def convert(self, masks, target_count):
        # Call the convert_masks_to_patches function
        patch_masks, h, w = convert_masks_to_patches(masks, target_count)
        print(f'Patch dimensions: {h}x{w}, Total patches: {h*w}')
        return (patch_masks, h, w)

class PreviewSubjectMask:
    RETURN_TYPES = ("MASK","LATENT",)
    RETURN_NAMES = ("mask","latent",)
    FUNCTION = "go"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT",),
                "subject": ("SUBJECT",),
            },
        }
    @classmethod
    def go(cls, latent, subject):
        # Iterate through subject[0] and gather all fg_masks
        fg_masks = []
        for s in subject:
            # Now s is a dict where keys are step ids
            for step_id, step_data in s.items():
                if "fg_mask" in step_data:
                    fg_masks.append(step_data["fg_mask"])
        
        # Stack all fg_masks into a single tensor
        if fg_masks:
            # Stack the masks to create a tensor with shape [B, H, W]
            combined_mask = torch.cat(fg_masks, dim=0)
        else:
            # If no fg_masks found, return empty tensor
            combined_mask = torch.empty(0)
            
        return (combined_mask, latent, )

class SetSubjectMask:
    RETURN_TYPES = ("SUBJECT",)
    RETURN_NAMES = ("subject",)
    FUNCTION = "set_mask"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mask": ("MASK",),
                "subject": ("SUBJECT",),
            }
        }

    @classmethod
    def set_mask(cls, mask: torch.Tensor, subject: list):
        """
        Set the global fg_mask for a subject based on an input mask.
        
        Args:
            mask: Input mask tensor with shape [B, H, W]
            subject: Subject list containing dictionaries with step data
            
        Returns:
            subject: Modified subject with global fg_mask set
        """
        # Get the dimensions of the input mask
        B, H, W = mask.shape
        
        # Calculate patch dimensions (divide by 16)
        h_patches = H // 16
        w_patches = W // 16
        
        # Reshape mask to [B, H//16, 16, W//16, 16] for processing 16x16 blocks
        mask_reshaped = mask.view(B, h_patches, 16, w_patches, 16)
        
        # Check if any pixel in each 16x16 block is non-zero
        # Take max over the 16x16 blocks, then convert to boolean
        patch_mask = mask_reshaped.max(dim=2)[0].max(dim=3)[0] > 0
        
        # Convert to boolean tensor
        patch_mask = patch_mask.bool()
        
        # Create a copy of the subject to avoid modifying the original
        new_subject = copy.deepcopy(subject)
        
        # Set the global fg_mask in subject[0]['fg_mask']
        # This represents the global mask of the subject, not for each step's fg_mask
        if len(new_subject) > 0:
            if 0 not in new_subject[0]:
                new_subject[0] = {}
            if 'fg_mask' not in new_subject[0][0]:
                new_subject[0][0]['fg_mask'] = patch_mask
            else:
                # If fg_mask already exists, replace it with our new global mask
                new_subject[0][0]['fg_mask'] = patch_mask
        
        return (new_subject,)

def hook_prepare_callback():
    flux_layers.attention = orig_attention
    chroma_layers.attention = orig_attention
    def prepare_callback(model, max_steps, x0_output_dict=None):
        if "prepare_callback" in model.model_options:
            return model.model_options["prepare_callback"](model, max_steps, x0_output_dict)
        else:
            return orig_prepare_callback(model, max_steps, x0_output_dict) 
    latent_preview.prepare_callback = prepare_callback

def get_interpolate_weight(weight, start_step, decay_step, end_step):
    steps = np.arange(0, end_step - decay_step)
    decay_weights = weight * 0.5 * (1 + np.cos(np.pi * steps / (end_step - decay_step)))
    decay_weights = decay_weights.tolist()
    constant_weights = [weight] * (decay_step - start_step)
    weight_list = constant_weights + decay_weights
    weight_dict = dict()
    for ind, interpolate_step in enumerate(range(start_step, end_step)):
        weight_dict[interpolate_step] = weight_list[ind]
    return weight_dict

def get_shared_fg_mask(id_fg_mask, curr_fg_mask, curr2id_argmax_indices, curr2id_valid_mask):
    curr_valid_mask = curr2id_valid_mask.flatten()
    id_fg_mask = id_fg_mask.flatten().to(curr2id_argmax_indices.device)
    curr_fg_mask = curr_fg_mask.flatten().to(curr2id_argmax_indices.device)
    rearrange_id_fg_mask = id_fg_mask[curr2id_argmax_indices]
    share_fg_mask = curr_fg_mask & rearrange_id_fg_mask & curr_valid_mask
    id_share_fg_indices = curr2id_argmax_indices[share_fg_mask]
    curr_share_fg_indices = torch.nonzero(share_fg_mask).squeeze()
    return id_share_fg_indices, curr_share_fg_indices

def get_fg_attn_weights(text_len, image_start, q, k, fg_condition_mask):
    attn_weights = torch.matmul(q[:, :, image_start:], k[:, :, :text_len].transpose(-2, -1)).to(torch.float16)
    attn_weights = attn_weights.softmax(-1).sum(1, keepdim=True)
    fg_condition_mask_gpu = fg_condition_mask
    bg_condition_mask_gpu = 1.0 - fg_condition_mask_gpu
    masks = torch.nn.functional.normalize(torch.stack([fg_condition_mask_gpu, bg_condition_mask_gpu], dim=1), p=1, dim=-1)
    attn_weights = (attn_weights * masks).sum(-1)
    return attn_weights

class GenSubject:
    RETURN_TYPES = ("MODEL","SUBJECT",)
    FUNCTION = "go"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "fg_condition_mask": ("MASK",),
                "kv_start_step":  ("INT",{"default":1,"min":0}),
                "kv_end_step":  ("INT",{"default":9,"min":0}),
                "x_start_step":  ("INT",{"default":1,"min":0}),
                "x_end_step":  ("INT",{"default":9,"min":0}),
                "mask_start_step": ("INT",{"default":8,"min":0}),
                "mask_end_step": ("INT",{"default":9,"min":0}),
            }
        }
    @classmethod
    def go(cls, model, fg_condition_mask, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step):
        
        subject = [{}]
        index = len(subject) - 1
        def prepare_callback(model, max_steps, x0_output_dict=None):
            nonlocal subject, fg_condition_mask
            subject[index] = {}
            sampled_step = -1
            sampling_step = -1
            layer = 0
            model_type = "flux"

            text_len = fg_condition_mask.shape[-1]
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
            fg_bg_mask = None
            fg_condition_mask = fg_condition_mask.to(device)
            def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
                """
                Implementation using PyTorch's scaled_dot_product_attention with custom masking.
                
                Args:
                    q: Query tensor of shape [B, H, L, D]
                    k: Key tensor of shape [B, H, L, D]
                    v: Value tensor of shape [B, H, L, D]
                    pe: Positional encoding tensor
                    attn_mask: Optional attention mask of shape [B, H, L, L]
                    
                Returns:
                    Output tensor of shape [B, L, H*D]
                """
                nonlocal sampling_step, sampled_step, layer, subject, text_len, fg_bg_mask
                if sampling_step != sampled_step + 1:
                    layer = 0
                    sampling_step = sampled_step + 1
                # Initialize step in subject dict if not exists
                if sampling_step not in subject[index]:
                    subject[index][sampling_step] = {"layers": {}}
                # Store k, v, and x in a dict with layer id as key
                if sampling_step >= kv_start_step and sampling_step < kv_end_step:
                    subject[index][sampling_step]["layers"][layer] = {"k": k.cpu(), "v": v.cpu()}
                # Apply rotary positional encoding for Flux models
                if pe is not None:
                    q, k = flux_math.apply_rope(q, k, pe)
                
                if sampling_step >= mask_start_step and sampling_step < mask_end_step:
                    if fg_bg_mask is None:
                        fg_bg_mask = get_fg_attn_weights(text_len, text_len, q, k, fg_condition_mask)
                    else:
                        fg_bg_mask += get_fg_attn_weights(text_len, text_len, q, k, fg_condition_mask)

                x = F.scaled_dot_product_attention(
                    q, k, v, attn_mask=mask
                )
                # Reshape back to the expected output format
                x = rearrange(x, "B H L D -> B L (H D)")
                # Store attention result x along with k and v
                if sampling_step >= x_start_step and sampling_step < x_end_step:
                    if not layer in subject[index][sampling_step]["layers"]:
                        subject[index][sampling_step]["layers"][layer] = {}
                    subject[index][sampling_step]["layers"][layer]["x"] = F.normalize(x[:, text_len:,:], dim=-1).transpose(-2, -1).cpu()
                layer += 1
                return x

            # Create the attention override function
            def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
                if transformer_options["cond_or_uncond"][-1]:#q.shape[2] != text_len + (w * h)
                    return orig_attention(q, k, v, pe)
                else:
                    return masked_attention(q, k, v, pe, mask)
            
            # Apply the attention override based on model type
            overrider = attention_override

            if model.model.__class__.__name__ == "Chroma":
                model_type = "chroma"
            
            print(f'GenSubject Detected model type: {model_type}')

            if model_type == "flux":
                flux_layers.attention = overrider
            else:  # chroma
                chroma_layers.attention = overrider            
            origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict)
            def callback_change_sampled_step(step, x0, x, total_steps):
                nonlocal max_steps
                nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model
                sampled_step = step
                if sampling_step >= mask_start_step and sampling_step < mask_end_step:
                    B,L,H,W = x0.shape
                    w = W // model.model.diffusion_model.patch_size
                    h = H // model.model.diffusion_model.patch_size
                    bitmask = fg_bg_mask[:,0] >= fg_bg_mask[:,1]
                    bitmask = rearrange(bitmask, "b (h w) -> b h w", h=h, w=w)
                    fg_bg_mask = None
                    # Store fg_mask in a dictionary with step id as key
                    if sampling_step not in subject[index]:
                        subject[index][sampling_step] = {"layers": {}}
                    subject[index][sampling_step]["fg_mask"] = bitmask.cpu()
                origin_callback(step, x0, x, total_steps)
                #if step == total_steps - 1:
                #    latent_preview.prepare_callback = orig_prepare_callback
            return callback_change_sampled_step
        model = copy.copy(model)
        model.model_options = copy.copy(model.model_options)
        model.model_options["prepare_callback"] = prepare_callback
        hook_prepare_callback()        
        return (model, subject,)

class PreRun:
    RETURN_TYPES = ("MODEL","CROSSSIM",)
    FUNCTION = "go"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "subject": ("SUBJECT",),
                "fg_condition_mask": ("MASK",),
                "cross_sim_step": ("INT",),
                "sim_thredhold": ("FLOAT",{"default":0.5,"min":0,"max":2}),
            },
            "optional":{
                "latent": ("LATENT",),
            }

        }
    @classmethod
    def go(cls, model, subject, fg_condition_mask, cross_sim_step,sim_thredhold,latent):
        cross_sim = {}
        
        def prepare_callback(model, max_steps, x0_output_dict=None):
            nonlocal fg_condition_mask, cross_sim_step
            sampled_step = -1
            sampling_step = -1
            layer = 0
            model_type = "flux"

            text_len = fg_condition_mask.shape[-1]
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
            fg_bg_mask = None
            accum_sim = None            
            fg_condition_mask = fg_condition_mask.to(device)
            def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
                """
                Implementation using PyTorch's scaled_dot_product_attention with custom masking.
                
                Args:
                    q: Query tensor of shape [B, H, L, D]
                    k: Key tensor of shape [B, H, L, D]
                    v: Value tensor of shape [B, H, L, D]
                    pe: Positional encoding tensor
                    attn_mask: Optional attention mask of shape [B, H, L, L]
                    
                Returns:
                    Output tensor of shape [B, L, H*D]
                """
                nonlocal sampling_step, sampled_step, layer, subject, text_len, accum_sim, fg_bg_mask
                if sampling_step != sampled_step + 1:
                    layer = 0
                    sampling_step = sampled_step + 1
                
                if pe is not None:
                    q, k = flux_math.apply_rope(q, k, pe)

                if sampling_step == cross_sim_step:
                    accum_mask = []
                    for i in range(fg_condition_mask.shape[0]):
                        accum_mask.append(get_fg_attn_weights(text_len, text_len, q, k, fg_condition_mask[i:i+1]))
                    accum_mask = torch.cat(accum_mask)
                    if fg_bg_mask is None:
                        fg_bg_mask = accum_mask
                    else:
                        fg_bg_mask += accum_mask
                x = F.scaled_dot_product_attention(
                    q, k, v, attn_mask=mask
                )
                # Reshape back to the expected output format
                x = rearrange(x, "B H L D -> B L (H D)")
                # Store attention result x along with k and v
                if sampling_step == cross_sim_step:
                    all_cross_sim = []
                    for s in subject:
                        all_cross_sim.append(torch.matmul(F.normalize(x[:, text_len:, :], dim=-1), s[sampling_step]["layers"][layer]["x"].to(device = device, dtype = attn_dtype)))
                    this_cross_sim = torch.cat(all_cross_sim)
                    if accum_sim is None:
                        accum_sim = this_cross_sim
                    else:
                        accum_sim += this_cross_sim
                layer += 1
                return x

            # Create the attention override function
            def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
                if transformer_options["cond_or_uncond"][-1]:#q.shape[2] != text_len + (w * h)
                    return orig_attention(q, k, v, pe)
                else:
                    return masked_attention(q, k, v, pe, mask)
            
            # Apply the attention override based on model type
            overrider = attention_override

            if model.model.__class__.__name__ == "Chroma":
                model_type = "chroma"
            
            print(f'GenSubject Detected model type: {model_type}')

            if model_type == "flux":
                flux_layers.attention = overrider
            else:  # chroma
                chroma_layers.attention = overrider            
            origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict)
            def callback_change_sampled_step(step, x0, x, total_steps):     
                nonlocal sampling_step, sampled_step, layer, model, cross_sim, accum_sim, cross_sim_step,fg_bg_mask
                sampled_step = step
                if sampling_step == cross_sim_step:
                    B,L,H,W = x0.shape
                    w = W // model.model.diffusion_model.patch_size
                    h = H // model.model.diffusion_model.patch_size
                    bitmask = fg_bg_mask[:,0] >= fg_bg_mask[:,1]
                    bitmask = rearrange(bitmask, "b (h w) -> b h w", h=h, w=w)
                    fg_bg_mask = None
                    value = accum_sim / layer
                    max_sim, argmax_indices = torch.max(value, dim=-1)
                    all_id_fg_inds = []
                    all_curr_fg_inds = []
                    for i in range(len(fg_condition_mask)):
                        if 0 in subject[i] and "fg_mask" in subject[i][0]:
                            fg_mask = subject[i][0]["fg_mask"][0]
                        else:
                            fg_mask = subject[i][sampling_step]["fg_mask"][0]
                        id_fg_inds, curr_fg_inds = get_shared_fg_mask(
                            fg_mask.to(device), 
                            bitmask[i], 
                            argmax_indices[i],
                            max_sim[i]>sim_thredhold
                        )
                        all_id_fg_inds.append(id_fg_inds)
                        all_curr_fg_inds.append(curr_fg_inds)

                    cross_sim["id_fg_inds"] = torch.stack(all_id_fg_inds).cpu()
                    cross_sim["curr_fg_inds"] = torch.stack(all_curr_fg_inds).cpu()
                    cross_sim["max_sim"] = max_sim.cpu()
                    cross_sim["argmax_indices"] = argmax_indices.cpu()
                    cross_sim["debug_mask"] = value.to(dtype=torch.float32).cpu()
                    cross_sim["fg_mask"] = bitmask.cpu()
            return callback_change_sampled_step
        
        model = copy.copy(model)
        model.model_options = copy.copy(model.model_options)
        model.model_options["prepare_callback"] = prepare_callback
        hook_prepare_callback()        
        return (model, cross_sim,)


class CrossSimMaskExtract:
    RETURN_TYPES = ("MASK", "MASK")
    RETURN_NAMES = ("cross_mask", "fg_mask")
    FUNCTION = "extract"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "cross_sim": ("CROSSSIM",),
                "latent": ("LATENT",),
            }
        }

    def extract(self, cross_sim, latent):
        # Extract fg_mask directly from cross_sim
        fg_mask = cross_sim.get("fg_mask", None)
        
        # Create cross_mask from max_sim or debug_mask
        # Using debug_mask as it appears to contain the cross attention similarity data
        cross_mask = cross_sim.get("debug_mask", None)
        
        # If debug_mask is not available, try to create from max_sim
        #if cross_mask is None and "max_sim" in cross_sim:
        #    # max_sim is likely a tensor with similarity values
        #    max_sim = cross_sim["max_sim"]
        #    # Normalize to 0-1 range for mask
        #    if max_sim is not None:
        #        cross_mask = (max_sim - max_sim.min()) / (max_sim.max() - max_sim.min() + 1e-8)
        
        # Handle case where neither is available
        #if fg_mask is None:
        #    # Create an empty mask if fg_mask is not available
        #    fg_mask = torch.zeros(1, 1, 1, dtype=torch.float32)
            
        if cross_mask is None:
            # Create an empty mask if cross_mask is not available
            cross_mask = torch.zeros(1, 1, 1, dtype=torch.float32)
            
        return (cross_mask, fg_mask)


class GenConsistent:
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "go"
    CATEGORY = "chara_consist"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "subject": ("SUBJECT",),
                "cross_sim": ("CROSSSIM",),
                "fg_condition_mask": ("MASK",),
                "kv_start_step":  ("INT",{"default":1,"min":0}),
                "kv_end_step":  ("INT",{"default":5,"min":0}),
            },
            "optional": {
                "latent":("LATENT",)
            }
        }
    @classmethod
    def go(cls, model, subject, cross_sim, fg_condition_mask, kv_start_step, kv_end_step, latent):
        
        index = len(subject) - 1
        def prepare_callback(model, max_steps, x0_output_dict=None):
            nonlocal subject, fg_condition_mask
            sampled_step = -1
            sampling_step = -1
            layer = 0
            model_type = "flux"

            text_len = fg_condition_mask.shape[-1]
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
            fg_bg_mask = None
            fg_condition_mask = fg_condition_mask.to(device)
            def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor) -> Tensor:
                """
                Implementation using PyTorch's scaled_dot_product_attention with custom masking.
                
                Args:
                    q: Query tensor of shape [B, H, L, D]
                    k: Key tensor of shape [B, H, L, D]
                    v: Value tensor of shape [B, H, L, D]
                    pe: Positional encoding tensor
                    attn_mask: Optional attention mask of shape [B, H, L, L]
                    
                Returns:
                    Output tensor of shape [B, L, H*D]
                """
                nonlocal sampling_step, sampled_step, layer, subject, text_len, fg_bg_mask
                if sampling_step != sampled_step + 1:
                    layer = 0
                    sampling_step = sampled_step + 1
                
                # Apply rotary positional encoding for Flux models
                if pe is not None:
                    q, k = flux_math.apply_rope(q, k, pe)
                # Store k, v, and x in a dict with layer id as key
                if sampling_step >= kv_start_step and sampling_step < kv_end_step:
                    # For all other steps, we may apply the key-value expansion algorithm
                    # Get cross-sim data for this layer
                    id_fg_inds = cross_sim.get("id_fg_inds", None)
                    curr_fg_inds = cross_sim.get("curr_fg_inds", None)
                    
                    # Process all subjects, not just the last one (index = len(subject) - 1)
                    for subj_idx in range(len(subject)):
                        # Check if we have valid cross-sim data and subject data to work with
                        if (id_fg_inds is not None and curr_fg_inds is not None and 
                            len(id_fg_inds) > 0 and len(curr_fg_inds) > 0 and
                            sampling_step in subject[subj_idx] and 
                            layer in subject[subj_idx][sampling_step]["layers"] and
                            "k" in subject[subj_idx][sampling_step]["layers"][layer] and 
                            "v" in subject[subj_idx][sampling_step]["layers"][layer]):
                            
                            ori_saved_key = subject[subj_idx][sampling_step]["layers"][layer]["k"].to(k.device)
                            ori_saved_value = subject[subj_idx][sampling_step]["layers"][layer]["v"].to(v.device)
                            
                            # Ensure we have valid indices within tensor bounds
                            # Handle batch dimension properly - cross_sim tensors have [B, ...] form where B = len(subject)
                            if subj_idx < id_fg_inds.shape[0] and subj_idx < curr_fg_inds.shape[0]:
                                subj_id_fg_inds = id_fg_inds[subj_idx] if id_fg_inds.dim() > 1 else id_fg_inds
                                subj_curr_fg_inds = curr_fg_inds[subj_idx] if curr_fg_inds.dim() > 1 else curr_fg_inds
                                
                                valid_id_inds = subj_id_fg_inds[subj_id_fg_inds < ori_saved_key.shape[2]]
                                valid_curr_inds = subj_curr_fg_inds[subj_curr_fg_inds < k.shape[2]]
                                
                                if len(valid_id_inds) > 0 and len(valid_curr_inds) > 0:
                                    # Apply key-value expansion based on cross-sim data
                                    # Extract the saved key-value pairs using valid indices
                                    saved_key_fg = ori_saved_key[:, :, valid_id_inds]
                                    saved_value_fg = ori_saved_value[:, :, valid_id_inds]
                                    
                                    # Apply rotary positional encoding to saved keys (fix missing apply_rope)
                                    if pe is not None:
                                        # Extract positional encoding for the saved key positions
                                        # This mimics the approach in attention_processor_characonsist.py
                                        image_rotary_emb_fg = pe[:,:,valid_curr_inds]
                                        saved_key_fg = flux_math.apply_rope1(saved_key_fg, image_rotary_emb_fg)
                                    
                                    # Concatenate the expanded key-value pairs
                                    k = torch.cat([k, saved_key_fg], dim=2)
                                    v = torch.cat([v, saved_value_fg], dim=2)
                
                x = F.scaled_dot_product_attention(
                    q, k, v, attn_mask=mask
                )
                # Reshape back to the expected output format
                x = rearrange(x, "B H L D -> B L (H D)")

                layer += 1
                return x

            # Create the attention override function
            def attention_override(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask: Tensor, transformer_options) -> Tensor:
                if transformer_options["cond_or_uncond"][-1]:#q.shape[2] != text_len + (w * h)
                    return orig_attention(q, k, v, pe)
                else:
                    return masked_attention(q, k, v, pe, mask)
            
            # Apply the attention override based on model type
            overrider = attention_override

            if model.model.__class__.__name__ == "Chroma":
                model_type = "chroma"
            
            print(f'GenSubject Detected model type: {model_type}')

            if model_type == "flux":
                flux_layers.attention = overrider
            else:  # chroma
                chroma_layers.attention = overrider            
            origin_callback = orig_prepare_callback(model, max_steps, x0_output_dict)
            def callback_change_sampled_step(step, x0, x, total_steps):
                nonlocal max_steps
                nonlocal sampling_step, sampled_step, layer, fg_bg_mask, model
                sampled_step = step
                origin_callback(step, x0, x, total_steps)
            return callback_change_sampled_step
        model = copy.copy(model)
        model.model_options = copy.copy(model.model_options)
        model.model_options["prepare_callback"] = prepare_callback
        hook_prepare_callback()        
        return (model, subject,)