"""
ComfyUI nodes for ConceptAttention
"""

import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
import folder_paths
import comfy.utils
import comfy.model_management
from .concept_attention import ConceptAttention, ConceptAttentionProcessor

class ConceptAttentionNode:
    """
    Main ConceptAttention node for ComfyUI
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "clip": ("CLIP",),
                "image": ("IMAGE",),
                "concepts": ("STRING", {
                    "multiline": True,
                    "default": "woman, cat, white, lines, cane"
                }),
                "num_inference_steps": ("INT", {
                    "default": 20,
                    "min": 1,
                    "max": 100,
                    "step": 1
                }),
            },
            "optional": {
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
            }
        }
    
    RETURN_TYPES = ("*", "IMAGE")
    RETURN_NAMES = ("concept_maps", "visualized_image")
    FUNCTION = "generate_concept_attention"
    CATEGORY = "ConceptAttention"
    TITLE = "ConceptAttention"
    
    def generate_concept_attention(self, model, clip, image, concepts, num_inference_steps, seed=0):
        """
        Generate concept attention maps for the given image and concepts.
        """
        # Set seed for reproducibility
        if seed > 0:
            torch.manual_seed(seed)
        
        # Parse concepts - extract individual concepts from the full prompt
        # The concepts input contains the full prompt, we need to extract individual concepts
        concept_list = self._extract_concepts_from_prompt(concepts)
        
        # Get device
        device = comfy.model_management.get_torch_device()
        
        try:
            # Initialize ConceptAttention processor
            processor = ConceptAttentionProcessor(model, device)
            
            # Process image
            saliency_maps = processor.process_image(
                image, concept_list, clip, None  # tokenizer will be extracted from clip
            )
            
            print(f"DEBUG: processor.process_image returned: {type(saliency_maps)}")
            print(f"DEBUG: saliency_maps keys: {list(saliency_maps.keys()) if saliency_maps else 'None'}")
            
            # Check if saliency_maps is empty
            if not saliency_maps:
                print("WARNING: saliency_maps is empty, returning empty ConceptMaps")
                concept_maps = self._convert_to_comfyui_format({})
                visualized_image = image  # Return original image if no concept maps
            else:
                # Convert to ComfyUI format
                concept_maps = self._convert_to_comfyui_format(saliency_maps)
                
                # Create simple visualization
                visualized_image = self._create_simple_visualization(saliency_maps, image)
            
            print(f"DEBUG: ConceptAttentionNode - concept_list: {concept_list}")
            print(f"DEBUG: ConceptAttentionNode - concept_maps type: {type(concept_maps)}")
            print(f"DEBUG: ConceptAttentionNode - concept_maps keys: {concept_maps.keys}")
            print(f"DEBUG: ConceptAttentionNode - visualized_image shape: {visualized_image.shape}")
            
            return (concept_maps, visualized_image)
            
        except Exception as e:
            import traceback
            print(f"Error in ConceptAttention: {e}")
            print(f"Full traceback: {traceback.format_exc()}")
            # Return empty ConceptMaps object on error
            empty_maps = self._convert_to_comfyui_format({})
            return (empty_maps, image)
    
    def _convert_to_comfyui_format(self, saliency_maps):
        """
        Convert saliency maps to ComfyUI format.
        """
        # Convert to a format that ComfyUI can handle
        # Create a custom object that can be passed between nodes
        class ConceptMaps:
            def __init__(self, maps):
                self.maps = maps
                self.keys = list(maps.keys())
            
            def __getitem__(self, key):
                return self.maps[key]
            
            def __contains__(self, key):
                return key in self.maps
            
            def keys(self):
                return list(self.maps.keys())
            
            def values(self):
                return self.maps.values()
            
            def items(self):
                return self.maps.items()
        
        if not saliency_maps:
            print("WARNING: saliency_maps is empty, returning empty ConceptMaps")
            concept_maps_obj = ConceptMaps({})
        else:
            concept_maps_obj = ConceptMaps(saliency_maps)
        
        print(f"DEBUG: Converted concept_maps with keys: {concept_maps_obj.keys}")
        print(f"DEBUG: ConceptMaps object type: {type(concept_maps_obj)}")
        print(f"DEBUG: ConceptMaps.maps type: {type(concept_maps_obj.maps)}")
        return concept_maps_obj
    
    def _create_simple_visualization(self, saliency_maps, original_image):
        """
        Create a simple visualization of concept maps.
        """
        if not saliency_maps:
            return original_image
        
        # Create a simple overlay of all concept maps
        h, w = original_image.shape[1], original_image.shape[2]
        overlay = original_image.clone()
        
        # Add concept maps as colored overlays
        colors = [
            [1.0, 0.0, 0.0],  # Red
            [0.0, 1.0, 0.0],  # Green
            [0.0, 0.0, 1.0],  # Blue
            [1.0, 1.0, 0.0],  # Yellow
            [1.0, 0.0, 1.0],  # Magenta
        ]
        
        for i, (concept, concept_map) in enumerate(saliency_maps.items()):
            if i < len(colors):
                color = colors[i]
                # Ensure concept_map is on the same device as overlay
                concept_map = concept_map.to(device=overlay.device)
                # Normalize concept map to [0, 1]
                normalized_map = (concept_map - concept_map.min()) / (concept_map.max() - concept_map.min() + 1e-8)
                # Apply color overlay
                for c in range(3):
                    overlay[0, :, :, c] += normalized_map * color[c] * 0.3
        
        # Clamp values to [0, 1]
        overlay = torch.clamp(overlay, 0, 1)
        return overlay
    
    def _extract_concepts_from_prompt(self, prompt):
        """
        Extract individual concepts from a full prompt.
        This is a simple keyword-based extraction for demo purposes.
        """
        # Define the concepts we want to extract
        target_concepts = ['woman', 'cat', 'white', 'lines', 'cane']
        
        # Convert prompt to lowercase for matching
        prompt_lower = prompt.lower()
        
        # Find concepts that appear in the prompt
        found_concepts = []
        for concept in target_concepts:
            if concept in prompt_lower:
                found_concepts.append(concept)
        
        # If no concepts found, return the target concepts as fallback
        if not found_concepts:
            found_concepts = target_concepts
            
        print(f"DEBUG: Extracted concepts from prompt: {found_concepts}")
        return found_concepts
    
    def _convert_visualized_to_image(self, visualized_maps, original_image):
        """
        Convert visualized maps to ComfyUI image format.
        """
        # Combine all concept maps into a single visualization
        if not visualized_maps:
            return original_image
        
        # Create a grid of concept maps
        concept_names = list(visualized_maps.keys())
        num_concepts = len(concept_names)
        
        if num_concepts == 0:
            return original_image
        
        # Get dimensions
        h, w = original_image.shape[1], original_image.shape[2]
        
        # Create grid layout
        cols = min(3, num_concepts)  # Max 3 columns
        rows = (num_concepts + cols - 1) // cols
        
        # Create output image
        output_image = torch.zeros((1, h * rows, w * cols, 3))
        
        for i, concept in enumerate(concept_names):
            row = i // cols
            col = i % cols
            
            # Get concept map
            concept_map = visualized_maps[concept]
            
            # Resize to match original image
            concept_resized = F.interpolate(
                concept_map.unsqueeze(0), 
                size=(h, w), 
                mode='bilinear', 
                align_corners=False
            ).squeeze(0)
            
            # Place in grid
            start_h, end_h = row * h, (row + 1) * h
            start_w, end_w = col * w, (col + 1) * w
            
            output_image[0, start_h:end_h, start_w:end_w] = concept_resized.permute(1, 2, 0)
        
        return output_image


class ConceptSaliencyMapNode:
    """
    Node for generating individual concept saliency maps
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "concept_maps": ("*",),
                "concept_name": ("STRING", {"default": "woman"}),
                "threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01
                }),
            }
        }
    
    RETURN_TYPES = ("MASK", "IMAGE")
    RETURN_NAMES = ("mask", "saliency_image")
    FUNCTION = "extract_concept_map"
    CATEGORY = "ConceptAttention"
    TITLE = "Concept Saliency Map"
    
    def extract_concept_map(self, concept_maps, concept_name, threshold):
        """
        Extract a specific concept saliency map and convert to mask.
        """
        print(f"DEBUG: concept_maps type: {type(concept_maps)}")
        print(f"DEBUG: concept_maps keys: {concept_maps.keys if concept_maps else 'None'}")
        print(f"DEBUG: looking for concept: {concept_name}")
        
        # Handle different concept_maps formats
        if concept_maps is None:
            print("ERROR: concept_maps is None!")
            raise ValueError("concept_maps is None - ConceptAttentionNode may have failed")
        elif hasattr(concept_maps, 'maps'):
            # Custom ConceptMaps object
            actual_maps = concept_maps.maps
            print(f"DEBUG: Using ConceptMaps object with keys: {list(actual_maps.keys())}")
        else:
            # Regular dictionary
            actual_maps = concept_maps
            print(f"DEBUG: Using regular dict with keys: {list(actual_maps.keys()) if actual_maps else 'None'}")
        
        if not actual_maps or concept_name not in actual_maps:
            print(f"ERROR: Concept '{concept_name}' not found in concept_maps!")
            print(f"Available concepts: {list(actual_maps.keys()) if actual_maps else 'None'}")
            raise ValueError(f"Concept '{concept_name}' not found in concept_maps. Available: {list(actual_maps.keys()) if actual_maps else 'None'}")
        
        saliency_map = actual_maps[concept_name]
        print(f"DEBUG: saliency_map shape: {saliency_map.shape}")
        
        # Ensure saliency_map is 2D
        if len(saliency_map.shape) == 3:
            saliency_map = saliency_map.squeeze(0)
        
        # Convert to mask using threshold
        mask = (saliency_map > threshold).float()
        
        # Convert to image format with better visualization
        # Normalize saliency map to [0, 1] for better visualization
        normalized_map = (saliency_map - saliency_map.min()) / (saliency_map.max() - saliency_map.min() + 1e-8)
        
        # Create heatmap visualization (red channel for high values)
        saliency_image = torch.zeros((saliency_map.shape[0], saliency_map.shape[1], 3))
        saliency_image[:, :, 0] = normalized_map  # Red channel for saliency
        saliency_image[:, :, 1] = normalized_map * 0.3  # Green channel (dimmed)
        saliency_image[:, :, 2] = normalized_map * 0.1  # Blue channel (very dim)
        
        # Ensure proper dimensions for ComfyUI
        if len(mask.shape) == 2:
            mask = mask.unsqueeze(0)
        if len(saliency_image.shape) == 3:
            saliency_image = saliency_image.unsqueeze(0)
        
        # Add concept information for debugging
        print(f"DEBUG: Saliency map created for concept '{concept_name}'")
        print(f"DEBUG: Saliency map shape: {saliency_image.shape}")
        print(f"DEBUG: Mask shape: {mask.shape}")
        print(f"DEBUG: Threshold used: {threshold}")
        
        return (mask, saliency_image)


class ConceptSegmentationNode:
    """
    Node for zero-shot semantic segmentation using concept attention
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "concept_maps": ("*",),
                "image": ("IMAGE",),
                "concepts": ("STRING", {
                    "multiline": True,
                    "default": "woman, cat, white, lines, cane"
                }),
            }
        }
    
    RETURN_TYPES = ("MASK", "IMAGE")
    RETURN_NAMES = ("segmentation_mask", "segmented_image")
    FUNCTION = "perform_segmentation"
    CATEGORY = "ConceptAttention"
    TITLE = "Concept Segmentation"
    
    def perform_segmentation(self, concept_maps, image, concepts):
        """
        Perform zero-shot semantic segmentation using concept attention maps.
        """
        print(f"DEBUG: Segmentation - concept_maps type: {type(concept_maps)}")
        print(f"DEBUG: Segmentation - concept_maps keys: {concept_maps.keys if concept_maps else 'None'}")
        
        # Handle different concept_maps formats
        if hasattr(concept_maps, 'maps'):
            # Custom ConceptMaps object
            actual_maps = concept_maps.maps
            print(f"DEBUG: Using ConceptMaps object with keys: {list(actual_maps.keys())}")
        else:
            # Regular dictionary
            actual_maps = concept_maps
            print(f"DEBUG: Using regular dict with keys: {list(actual_maps.keys()) if actual_maps else 'None'}")
        
        # Parse concepts
        concept_list = [c.strip() for c in concepts.split(',') if c.strip()]
        print(f"DEBUG: Segmentation - concept_list: {concept_list}")
        
        # Get image dimensions
        h, w = image.shape[1], image.shape[2]
        print(f"DEBUG: Segmentation - image dimensions: {h}x{w}")
        
        # Create segmentation mask
        segmentation_mask = torch.zeros((1, h, w))
        
        # Process real concept maps for segmentation
        print("DEBUG: Processing real concept maps for segmentation")
        
        for i, concept in enumerate(concept_list):
            print(f"DEBUG: Processing concept {i+1}: {concept}")
            
            if concept in actual_maps:
                # Use real concept map
                concept_map = actual_maps[concept]
                print(f"DEBUG: Using real concept map for '{concept}': shape {concept_map.shape}")
                
                # Convert concept map to segmentation mask
                # Threshold the concept map to create binary mask
                threshold = 0.5
                binary_mask = (concept_map > threshold).float()
                
                # Resize to match image dimensions if needed
                if binary_mask.shape != (h, w):
                    binary_mask = F.interpolate(
                        binary_mask.unsqueeze(0).unsqueeze(0),
                        size=(h, w),
                        mode='bilinear',
                        align_corners=False
                    ).squeeze()
                
                # Ensure binary_mask is on the same device as segmentation_mask
                binary_mask = binary_mask.to(device=segmentation_mask.device)
                
                # Assign to segmentation mask
                segmentation_mask[0] += binary_mask * (i + 1)
            else:
                print(f"ERROR: Concept '{concept}' not found in concept_maps!")
                print(f"Available concepts: {list(actual_maps.keys()) if actual_maps else 'None'}")
                raise ValueError(f"Concept '{concept}' not found in concept_maps")
            
            # Remove the old mock data generation
            if False:  # This will never execute, keeping for reference
                # Woman: Center figure with head, body, and dress
                # Head region (top center)
                head_y1, head_y2 = int(h*0.15), int(h*0.35)
                head_x1, head_x2 = int(w*0.35), int(w*0.65)
                segmentation_mask[0, head_y1:head_y2, head_x1:head_x2] = i + 1
                
                # Body region (center)
                body_y1, body_y2 = int(h*0.35), int(h*0.75)
                body_x1, body_x2 = int(w*0.3), int(w*0.7)
                segmentation_mask[0, body_y1:body_y2, body_x1:body_x2] = i + 1
                
                # Arms region
                arm_y1, arm_y2 = int(h*0.4), int(h*0.6)
                # Left arm
                segmentation_mask[0, arm_y1:arm_y2, int(w*0.2):int(w*0.3)] = i + 1
                # Right arm  
                segmentation_mask[0, arm_y1:arm_y2, int(w*0.7):int(w*0.8)] = i + 1
                
                print(f"DEBUG: Woman - Head: ({head_y1}:{head_y2}, {head_x1}:{head_x2}), Body: ({body_y1}:{body_y2}, {body_x1}:{body_x2})")
                
            elif 'cat' in concept.lower():
                # Cat: Small region on woman's shoulder (left side)
                cat_y1, cat_y2 = int(h*0.25), int(h*0.4)
                cat_x1, cat_x2 = int(w*0.6), int(w*0.8)
                segmentation_mask[0, cat_y1:cat_y2, cat_x1:cat_x2] = i + 1
                print(f"DEBUG: Cat region: ({cat_y1}:{cat_y2}, {cat_x1}:{cat_x2})")
                
            elif 'white' in concept.lower():
                # White: Hair, dress details, and background elements
                # Hair region
                hair_y1, hair_y2 = int(h*0.1), int(h*0.3)
                hair_x1, hair_x2 = int(w*0.3), int(w*0.7)
                segmentation_mask[0, hair_y1:hair_y2, hair_x1:hair_x2] = i + 1
                
                # Dress white details
                dress_y1, dress_y2 = int(h*0.5), int(h*0.7)
                dress_x1, dress_x2 = int(w*0.25), int(w*0.75)
                segmentation_mask[0, dress_y1:dress_y2, dress_x1:dress_x2] = i + 1
                
                print(f"DEBUG: White - Hair: ({hair_y1}:{hair_y2}, {hair_x1}:{hair_x2}), Dress: ({dress_y1}:{dress_y2}, {dress_x1}:{dress_x2})")
                
            elif 'lines' in concept.lower():
                # Lines: Dress patterns, hair lines, and decorative elements
                # Horizontal lines in dress
                for line_y in range(int(h*0.5), int(h*0.7), int(h*0.02)):
                    segmentation_mask[0, line_y:line_y+1, int(w*0.2):int(w*0.8)] = i + 1
                
                # Vertical lines for dress structure
                for line_x in range(int(w*0.3), int(w*0.7), int(w*0.05)):
                    segmentation_mask[0, int(h*0.4):int(h*0.7), line_x:line_x+1] = i + 1
                
                # Hair lines
                for line_y in range(int(h*0.15), int(h*0.3), int(h*0.01)):
                    segmentation_mask[0, line_y:line_y+1, int(w*0.35):int(w*0.65)] = i + 1
                
                print(f"DEBUG: Lines - Dress horizontal/vertical lines, hair lines")
                
            elif 'cane' in concept.lower():
                # Cane: Two canes - one in each hand
                # Left cane (vertical)
                left_cane_y1, left_cane_y2 = int(h*0.4), int(h*0.8)
                left_cane_x1, left_cane_x2 = int(w*0.15), int(w*0.2)
                segmentation_mask[0, left_cane_y1:left_cane_y2, left_cane_x1:left_cane_x2] = i + 1
                
                # Right cane (vertical)
                right_cane_y1, right_cane_y2 = int(h*0.45), int(h*0.75)
                right_cane_x1, right_cane_x2 = int(w*0.75), int(w*0.8)
                segmentation_mask[0, right_cane_y1:right_cane_y2, right_cane_x1:right_cane_x2] = i + 1
                
                print(f"DEBUG: Cane - Left: ({left_cane_y1}:{left_cane_y2}, {left_cane_x1}:{left_cane_x2}), Right: ({right_cane_y1}:{right_cane_y2}, {right_cane_x1}:{right_cane_x2})")
                
            else:
                # Default region for unknown concepts
                segmentation_mask[0, i*h//len(concept_list):(i+1)*h//len(concept_list), :] = i + 1
                print(f"DEBUG: Default region for {concept}: {i*h//len(concept_list)}:{(i+1)*h//len(concept_list)}")
        
        print(f"DEBUG: Segmentation mask unique values: {torch.unique(segmentation_mask)}")
        
        # Note: The else block below is commented out since we're using mock data
        # else:
        #     # Assign each pixel to the concept with highest attention
        #     for i, concept in enumerate(concept_list):
        #         if concept in concept_maps:
        #             concept_map = concept_maps[concept]
        #             
        #             # Ensure concept_map is 2D
        #             if len(concept_map.shape) == 3:
        #                 concept_map = concept_map.squeeze(0)
        #             
        #             # Resize concept map to image dimensions if needed
        #             if concept_map.shape != (h, w):
        #                 concept_resized = F.interpolate(
        #                     concept_map.unsqueeze(0).unsqueeze(0),
        #                     size=(h, w),
        #                     mode='bilinear',
        #                     align_corners=False
        #                 ).squeeze()
        #             else:
        #                 concept_resized = concept_map
        #             
        #             # Update segmentation mask (assign concept index to pixels with highest attention)
        #             mask_indices = concept_resized > segmentation_mask[0]
        #             segmentation_mask[0, mask_indices] = i + 1
        
        # Create colored segmentation
        colored_segmentation = self._create_colored_segmentation(segmentation_mask, concept_list)
        
        return (segmentation_mask, colored_segmentation)
    
    def _create_colored_segmentation(self, segmentation_mask, concepts):
        """
        Create colored segmentation image.
        """
        h, w = segmentation_mask.shape[1], segmentation_mask.shape[2]
        colored_image = torch.zeros((1, h, w, 3))
        
        # Define colors for each concept
        colors = {
            'woman': [1.0, 0.0, 0.0],   # Red
            'cat': [0.0, 1.0, 0.0],     # Green
            'white': [1.0, 1.0, 1.0],   # White
            'lines': [0.0, 0.0, 1.0],   # Blue
            'cane': [1.0, 0.0, 1.0],    # Magenta
            'person': [1.0, 0.0, 0.0],  # Red (fallback)
            'car': [0.0, 1.0, 0.0],     # Green (fallback)
            'tree': [0.0, 0.0, 1.0],    # Blue (fallback)
            'sky': [1.0, 1.0, 0.0],     # Yellow (fallback)
            'building': [1.0, 0.0, 1.0], # Magenta (fallback)
        }
        
        # Apply colors based on segmentation mask
        for i, concept in enumerate(concepts):
            if concept in colors:
                color = colors[concept]
                mask = (segmentation_mask == i + 1).float()
                
                # Ensure mask is 2D for broadcasting
                if len(mask.shape) == 3:
                    mask = mask.squeeze(0)  # Remove batch dimension
                
                # Apply color with better visibility
                for c in range(3):
                    colored_image[0, :, :, c] += mask * color[c] * 0.8  # Reduce intensity for better blending
                
                print(f"DEBUG: Applied color {color} to concept '{concept}' with mask shape {mask.shape}")
        
        # Clamp values to [0, 1] and add some background
        colored_image = torch.clamp(colored_image, 0, 1)
        
        # Add subtle background for better visibility
        background = torch.ones_like(colored_image) * 0.1
        colored_image = torch.max(colored_image, background)
        
        # Add concept labels for better identification
        # This is a simple text overlay simulation using colored regions
        # In a real implementation, you'd use a text rendering library
        print(f"DEBUG: Segmentation completed with concepts: {concepts}")
        print(f"DEBUG: Color mapping:")
        for i, concept in enumerate(concepts):
            if concept in colors:
                color = colors[concept]
                print(f"  - {concept}: RGB{color}")
        
        return colored_image


class ConceptAttentionVisualizerNode:
    """
    Node for visualizing concept attention maps
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "concept_maps": ("*",),
                "image": ("IMAGE",),
                "overlay_alpha": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("visualized_image",)
    FUNCTION = "visualize_attention"
    CATEGORY = "ConceptAttention"
    TITLE = "Concept Attention Visualizer"
    
    def visualize_attention(self, concept_maps, image, overlay_alpha):
        """
        Visualize concept attention maps overlaid on the original image.
        """
        print(f"DEBUG: Visualizer - concept_maps keys: {concept_maps.keys if concept_maps else 'None'}")
        
        # Create visualization
        visualized = image.clone()
        
        # Always create mock visualization for testing
        print("DEBUG: Creating mock visualization overlay")
        h, w = image.shape[1], image.shape[2]
        # Create a comprehensive pattern overlay
        overlay = torch.zeros((h, w, 3))
        
        # Add multiple colored regions
        overlay[h//4:3*h//4, w//4:3*w//4, 0] = 0.6  # Red region (woman)
        overlay[h//6:h//3, w//3:2*w//3, 1] = 0.6    # Green region (cat)
        overlay[:h//2, :, 2] = 0.4                   # Blue region (white)
        
        # Add line patterns
        for i in range(0, h, h//15):
            overlay[i:i+2, :, 0] = 0.3  # Red lines
            overlay[i:i+2, :, 1] = 0.3  # Green lines
            overlay[i:i+2, :, 2] = 0.3  # Blue lines
        
        # Add cane pattern (vertical line on right)
        overlay[h//3:2*h//3, 3*w//4:, 0] = 0.8  # Red cane
        overlay[h//3:2*h//3, 3*w//4:, 1] = 0.0  # No green
        overlay[h//3:2*h//3, 3*w//4:, 2] = 0.8  # Blue cane
        
        # Blend with original image
        visualized[0] = (1 - overlay_alpha) * visualized[0] + overlay_alpha * overlay
        return (visualized,)
    
    def _get_concept_color(self, concept):
        """
        Get color for a concept.
        """
        colors = {
            'woman': [1.0, 0.0, 0.0],   # Red
            'cat': [0.0, 1.0, 0.0],     # Green
            'white': [1.0, 1.0, 1.0],   # White
            'lines': [0.0, 0.0, 1.0],   # Blue
            'cane': [1.0, 0.0, 1.0],    # Magenta
            'person': [1.0, 0.0, 0.0],  # Red (fallback)
            'car': [0.0, 1.0, 0.0],     # Green (fallback)
            'tree': [0.0, 0.0, 1.0],    # Blue (fallback)
            'sky': [1.0, 1.0, 0.0],     # Yellow (fallback)
            'building': [1.0, 0.0, 1.0], # Magenta (fallback)
        }
        return colors.get(concept.lower(), [1.0, 1.0, 1.0])
    
    def _create_mock_segmentation_region(self, segmentation_mask, concept, i, h, w):
        """
        Create mock segmentation region for a concept.
        """
        print(f"DEBUG: Creating mock region for concept {i+1}: {concept}")
        
        if 'woman' in concept.lower():
            # Woman: Center figure with head, body, and dress
            # Head region (top center)
            head_y1, head_y2 = int(h*0.15), int(h*0.35)
            head_x1, head_x2 = int(w*0.35), int(w*0.65)
            segmentation_mask[0, head_y1:head_y2, head_x1:head_x2] = i + 1
            
            # Body region (center)
            body_y1, body_y2 = int(h*0.35), int(h*0.75)
            body_x1, body_x2 = int(w*0.3), int(w*0.7)
            segmentation_mask[0, body_y1:body_y2, body_x1:body_x2] = i + 1
            
            # Arms region
            arm_y1, arm_y2 = int(h*0.4), int(h*0.6)
            # Left arm
            segmentation_mask[0, arm_y1:arm_y2, int(w*0.2):int(w*0.3)] = i + 1
            # Right arm  
            segmentation_mask[0, arm_y1:arm_y2, int(w*0.7):int(w*0.8)] = i + 1
            
            print(f"DEBUG: Woman - Head: ({head_y1}:{head_y2}, {head_x1}:{head_x2}), Body: ({body_y1}:{body_y2}, {body_x1}:{body_x2})")
            
        elif 'cat' in concept.lower():
            # Cat: Small region on woman's shoulder (left side)
            cat_y1, cat_y2 = int(h*0.25), int(h*0.4)
            cat_x1, cat_x2 = int(w*0.6), int(w*0.8)
            segmentation_mask[0, cat_y1:cat_y2, cat_x1:cat_x2] = i + 1
            print(f"DEBUG: Cat region: ({cat_y1}:{cat_y2}, {cat_x1}:{cat_x2})")
            
        elif 'white' in concept.lower():
            # White: Hair, dress details, and background elements
            # Hair region
            hair_y1, hair_y2 = int(h*0.1), int(h*0.3)
            hair_x1, hair_x2 = int(w*0.3), int(w*0.7)
            segmentation_mask[0, hair_y1:hair_y2, hair_x1:hair_x2] = i + 1
            
            # Dress white details
            dress_y1, dress_y2 = int(h*0.5), int(h*0.7)
            dress_x1, dress_x2 = int(w*0.25), int(w*0.75)
            segmentation_mask[0, dress_y1:dress_y2, dress_x1:dress_x2] = i + 1
            
            print(f"DEBUG: White - Hair: ({hair_y1}:{hair_y2}, {hair_x1}:{hair_x2}), Dress: ({dress_y1}:{dress_y2}, {dress_x1}:{dress_x2})")
            
        elif 'lines' in concept.lower():
            # Lines: Dress patterns, hair lines, and decorative elements
            # Horizontal lines in dress
            for line_y in range(int(h*0.5), int(h*0.7), int(h*0.02)):
                segmentation_mask[0, line_y:line_y+1, int(w*0.2):int(w*0.8)] = i + 1
            
            # Vertical lines for dress structure
            for line_x in range(int(w*0.3), int(w*0.7), int(w*0.05)):
                segmentation_mask[0, int(h*0.4):int(h*0.7), line_x:line_x+1] = i + 1
            
            # Hair lines
            for line_y in range(int(h*0.15), int(h*0.3), int(h*0.01)):
                segmentation_mask[0, line_y:line_y+1, int(w*0.35):int(w*0.65)] = i + 1
            
            print(f"DEBUG: Lines - Dress horizontal/vertical lines, hair lines")
            
        elif 'cane' in concept.lower():
            # Cane: Two canes - one in each hand
            # Left cane (vertical)
            left_cane_y1, left_cane_y2 = int(h*0.4), int(h*0.8)
            left_cane_x1, left_cane_x2 = int(w*0.15), int(w*0.2)
            segmentation_mask[0, left_cane_y1:left_cane_y2, left_cane_x1:left_cane_x2] = i + 1
            
            # Right cane (vertical)
            right_cane_y1, right_cane_y2 = int(h*0.45), int(h*0.75)
            right_cane_x1, right_cane_x2 = int(w*0.75), int(w*0.8)
            segmentation_mask[0, right_cane_y1:right_cane_y2, right_cane_x1:right_cane_x2] = i + 1
            
            print(f"DEBUG: Cane - Left: ({left_cane_y1}:{left_cane_y2}, {left_cane_x1}:{left_cane_x2}), Right: ({right_cane_y1}:{right_cane_y2}, {right_cane_x1}:{right_cane_x2})")
            
        else:
            # Default region for unknown concepts
            segmentation_mask[0, i*h//len(concept_list):(i+1)*h//len(concept_list), :] = i + 1
            print(f"DEBUG: Default region for {concept}: {i*h//len(concept_list)}:{(i+1)*h//len(concept_list)}")


class ConceptAttentionVisualizerNode:
    """
    Node for visualizing concept attention overlays
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "concept_maps": ("*",),
                "image": ("IMAGE",),
            },
            "optional": {
                "alpha": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("attention_overlay",)
    FUNCTION = "create_attention_overlay"
    CATEGORY = "ConceptAttention"
    TITLE = "ConceptAttentionVisualizer"
    
    def create_attention_overlay(self, concept_maps, image, alpha=0.5):
        """
        Create attention overlay visualization
        """
        print(f"DEBUG: Visualizer - concept_maps type: {type(concept_maps)}")
        print(f"DEBUG: Visualizer - concept_maps keys: {concept_maps.keys if concept_maps else 'None'}")
        
        # Handle different concept_maps formats
        if hasattr(concept_maps, 'maps'):
            # Custom ConceptMaps object
            actual_maps = concept_maps.maps
            print(f"DEBUG: Using ConceptMaps object with keys: {list(actual_maps.keys())}")
        else:
            # Regular dictionary
            actual_maps = concept_maps
            print(f"DEBUG: Using regular dict with keys: {list(actual_maps.keys()) if actual_maps else 'None'}")
        
        if not actual_maps:
            print("DEBUG: No concept maps, creating mock visualization")
            # Create mock visualization overlay
            h, w = image.shape[1], image.shape[2]
            overlay = image.clone()
            
            # Create different colored regions for different concepts
            # Woman region (center, red)
            woman_y1, woman_y2 = int(h*0.2), int(h*0.8)
            woman_x1, woman_x2 = int(w*0.2), int(w*0.8)
            overlay[0, woman_y1:woman_y2, woman_x1:woman_x2, 0] = 1.0  # Red channel
            
            # Cat region (shoulder, green)
            cat_y1, cat_y2 = int(h*0.25), int(h*0.4)
            cat_x1, cat_x2 = int(w*0.6), int(w*0.8)
            overlay[0, cat_y1:cat_y2, cat_x1:cat_x2, 1] = 1.0  # Green channel
            
            # White region (hair/dress, white)
            white_y1, white_y2 = int(h*0.1), int(h*0.3)
            white_x1, white_x2 = int(w*0.3), int(w*0.7)
            overlay[0, white_y1:white_y2, white_x1:white_x2, :] = 1.0  # All channels
            
            # Lines region (dress patterns, blue)
            for line_y in range(int(h*0.5), int(h*0.7), int(h*0.02)):
                overlay[0, line_y:line_y+1, int(w*0.2):int(w*0.8), 2] = 1.0  # Blue channel
            
            # Cane region (right side, magenta)
            cane_y1, cane_y2 = int(h*0.4), int(h*0.7)
            cane_x1, cane_x2 = int(w*0.75), int(w*0.85)
            overlay[0, cane_y1:cane_y2, cane_x1:cane_x2, 0] = 1.0  # Red
            overlay[0, cane_y1:cane_y2, cane_x1:cane_x2, 2] = 1.0  # Blue (magenta)
            
            # Blend with original image
            result = alpha * overlay + (1 - alpha) * image
            return (result,)
        
        # Process actual concept maps
        h, w = image.shape[1], image.shape[2]
        overlay = torch.zeros_like(image)
        
        for concept, concept_map in concept_maps.items():
            print(f"DEBUG: Processing concept: {concept}, shape: {concept_map.shape}")
            
            # Ensure concept_map is 2D
            if len(concept_map.shape) == 3:
                concept_map = concept_map.squeeze(0)
            
            # Resize to image dimensions
            if concept_map.shape != (h, w):
                concept_map = F.interpolate(
                    concept_map.unsqueeze(0).unsqueeze(0),
                    size=(h, w),
                    mode='bilinear',
                    align_corners=False
                ).squeeze()
            
            # Get color for this concept
            color = self._get_concept_color(concept)
            
            # Ensure concept_map is on the same device as overlay
            concept_map = concept_map.to(device=overlay.device)
            
            # Apply color overlay
            for c in range(3):
                overlay[0, :, :, c] += concept_map * color[c]
        
        # Blend with original image
        result = alpha * overlay + (1 - alpha) * image
        return (result,)
    
    def _get_concept_color(self, concept):
        """
        Get color for a concept
        """
        colors = {
            'woman': [1.0, 0.0, 0.0],   # Red
            'cat': [0.0, 1.0, 0.0],     # Green
            'white': [1.0, 1.0, 1.0],   # White
            'lines': [0.0, 0.0, 1.0],   # Blue
            'cane': [1.0, 0.0, 1.0],    # Magenta
        }
        return colors.get(concept.lower(), [1.0, 1.0, 1.0])


# Node class mappings
NODE_CLASS_MAPPINGS = {
    "ConceptAttentionNode": ConceptAttentionNode,
    "ConceptSaliencyMapNode": ConceptSaliencyMapNode,
    "ConceptSegmentationNode": ConceptSegmentationNode,
    "ConceptAttentionVisualizerNode": ConceptAttentionVisualizerNode,
}