modules/infographics_generator/layout_system/hierarchical_optimizer.py

"""Hierarchical layout optimizer with extensible architecture."""

from typing import Dict, List, Optional, Tuple, Any
import numpy as np
from dataclasses import dataclass, field
from . import parameters as params
from .constraints import (
    ConstraintProcessor,
    RelativeSizeProcessor,
    PaddingProcessor,
    OrientationProcessor,
    OverlapProcessor,
    AlignmentProcessor,
    GapProcessor,
)
from .handlers import (
    NodeHandler,
    ImageNodeHandler,
    TextNodeHandler,
)
from .strategies import OptimizationStrategy, SDFOptimizationStrategy, RuleBasedLayoutStrategy
from .utils.parser import parse_layout_tree, LayoutNode
from .utils.composite import composite_nodes
from .utils.placeholder import create_placeholder_rectangle, create_placeholder_rounded_rectangle


@dataclass
class OptimizationConfig:
    """Optimization configuration."""
    strategy: OptimizationStrategy = field(default_factory=SDFOptimizationStrategy)
    constraint_processors: Dict[str, ConstraintProcessor] = field(default_factory=dict)
    node_handlers: Dict[str, NodeHandler] = field(default_factory=dict)
    weights: Dict[str, float] = field(default_factory=dict)
    optimization_params: Dict[str, Any] = field(default_factory=dict)
    placeholder_config: Dict[str, Any] = field(default_factory=dict)
    base_dir: Optional[str] = None
    device: Optional[str] = None
    debug: bool = False  # Enable debug mode: visualization and detailed logging
    use_rule_based: bool = True  # Enable rule-based layout for row/column nodes (faster)
    rule_based_types: List[str] = field(default_factory=lambda: ['row', 'column'])  # Types that use rule-based layout
    
    def __post_init__(self):
        """Initialize default values."""
        if not self.constraint_processors:
            self.constraint_processors = {
                "relative_size": RelativeSizeProcessor(),
                "padding": PaddingProcessor(),
                "orientation": OrientationProcessor(),
                "overlap": OverlapProcessor(),
                "gap": GapProcessor(),
            }
        
        if not self.node_handlers:
            self.node_handlers = {
                "image": ImageNodeHandler(),
                "chart": ImageNodeHandler(),
                "text": TextNodeHandler(),
            }
        
        if not self.weights:
            self.weights = {
                "w_similarity": params.W_SIMILARITY,
                "w_readability": params.W_READABILITY,
                "w_alignment_consistency": params.W_ALIGNMENT_CONSISTENCY,
                "w_alignment_similarity": params.W_ALIGNMENT_SIMILARITY,
                "w_proximity": params.W_PROXIMITY,
            }
        
        if not self.optimization_params:
            self.optimization_params = {
                "opt_res_list": params.OPT_RES_LIST,
                "outer_rounds": params.OUTER_ROUNDS,
                "inner_steps": params.INNER_STEPS,
                "lr": params.LEARNING_RATE,
            }


class HierarchicalOptimizer:
    """Hierarchical layout optimizer with extensible architecture."""
    
    def __init__(self, config: Optional[OptimizationConfig] = None):
        """Initialize optimizer.
        
        Args:
            config: Optimization configuration (uses default if None)
        """
        self.config = config or OptimizationConfig()
        self._register_default_processors()
        self._register_default_handlers()
    
    def _register_default_processors(self):
        """Register default constraint processors."""
        if not self.config.constraint_processors:
            self.config.constraint_processors = {
                "relative_size": RelativeSizeProcessor(),
                "padding": PaddingProcessor(),
                "orientation": OrientationProcessor(),
                "overlap": OverlapProcessor(),
            }
    
    def _register_default_handlers(self):
        """Register default node handlers."""
        if not self.config.node_handlers:
            self.config.node_handlers = {
                "image": ImageNodeHandler(),
                "chart": ImageNodeHandler(),
                "text": TextNodeHandler(),
            }
    
    def register_processor(self, processor: ConstraintProcessor):
        """Register a custom constraint processor.
        
        Args:
            processor: Constraint processor instance
        """
        # Register for all constraint types it can handle
        for constraint_type in ["relative_size", "padding", "orientation", "overlap", "gap", "alignment"]:
            if processor.can_handle(constraint_type):
                self.config.constraint_processors[constraint_type] = processor
    
    def register_handler(self, handler: NodeHandler):
        """Register a custom node handler.
        
        Args:
            handler: Node handler instance
        """
        # Register for all node types it can handle
        for node_type in ["image", "chart", "text", "shape", "layer", "column", "row"]:
            if handler.can_handle(node_type):
                self.config.node_handlers[node_type] = handler
    
    def set_strategy(self, strategy: OptimizationStrategy):
        """Set optimization strategy.
        
        Args:
            strategy: Optimization strategy instance
        """
        self.config.strategy = strategy
    
    def optimize_tree(self, tree_json: dict) -> Dict[str, Any]:
        """Optimize entire tree structure.
        
        Args:
            tree_json: JSON dictionary with layout tree structure
        
        Returns:
            Dictionary with optimization results for each node
        """
        # Parse tree
        root_node = parse_layout_tree(tree_json)
        # print("Root node: ", root_node)
        
        # Get container bbox from root
        root_bbox = (
            root_node.bbox.get("x", 0),
            root_node.bbox.get("y", 0),
            root_node.bbox.get("width", 1000),
            root_node.bbox.get("height", 1000),
        )
        
        # Optimize recursively from bottom up
        result = self._optimize_node(root_node, root_bbox, "root")
        
        return result
    
    def _optimize_node(self, node: LayoutNode, parent_bbox: Tuple[float, float, float, float], 
                       node_path: str = "root") -> Dict[str, Any]:
        """Optimize a single node and its children recursively.
        
        Args:
            node: Layout node to optimize
            parent_bbox: Parent container bounding box (x, y, w, h)
            node_path: Path string for this node (e.g., "root.child0.child1")
        
        Returns:
            Dictionary with optimization results
        """
        result = {
            "type": node.type,
            "bbox": node.bbox,
            "final_bbox": None,
            "image_path": getattr(node, "image_path", None),  # Preserve image_path for saving
        }
        
        # Check if node is a container (has children)
        if node.children:
            # Check if we can use rule-based layout (much faster for row/column)
            if self._can_use_rule_based_layout(node):
                return self._rule_based_layout(node, parent_bbox, node_path)
            
            # Container node: optimize children using SDF optimization
            # print(f"[HierarchicalOptimizer] Processing container node: type={node.type}, num_children={len(node.children)}, path={node_path}")
            child_results = []
            child_nodes_data = []
            
            # First, recursively optimize all children
            for i, child in enumerate(node.children):
                child_path = f"{node_path}.child{i}"
                child_result = self._optimize_node(child, parent_bbox, child_path)
                child_results.append(child_result)
            
            # Load child node data (masks, etc.)
            # Use final_bbox from child_result if available (for container nodes that have been optimized),
            # otherwise use initial bbox
            for i, child in enumerate(node.children):
                child_result = child_results[i]
                # Get bbox: use final_bbox from child_result if child is a container that has been optimized
                if child_result.get("final_bbox") is not None:
                    # Convert final_bbox tuple to dict format for consistency
                    final_bbox = child_result["final_bbox"]
                    if isinstance(final_bbox, (tuple, list)) and len(final_bbox) >= 4:
                        child_bbox = {
                            "x": final_bbox[0],
                            "y": final_bbox[1],
                            "width": final_bbox[2],
                            "height": final_bbox[3]
                        }
                    else:
                        child_bbox = child.bbox
                else:
                    child_bbox = child.bbox
                
                # Check if child has composite_mask (from previous optimization)
                # Use composite_mask if available, as it represents the actual shape
                mask = None
                metadata = {}
                if child_result.get("composite_mask") is not None:
                    mask = child_result["composite_mask"]
                    # print(f"[HierarchicalOptimizer] Using composite_mask for child {i} (type={child.type})")
                else:
                    handler = self._get_handler(child.type)
                    if handler:
                        mask, metadata = handler.load(child.__dict__, self.config.base_dir)
                    else:
                        # Fallback: use placeholder
                        bbox = child_bbox
                        width = bbox.get("width", 100) if isinstance(bbox, dict) else (bbox[2] if isinstance(bbox, (tuple, list)) else 100)
                        height = bbox.get("height", 100) if isinstance(bbox, dict) else (bbox[3] if isinstance(bbox, (tuple, list)) else 100)
                        # For container nodes (layer/column/row), create a rounded rectangle instead of solid rectangle
                        if child.type in ["layer", "column", "row"]:
                            try:
                                _, mask = create_placeholder_rounded_rectangle(width, height)
                            except Exception:
                                # Fallback to rectangle if rounded rectangle is unavailable
                                _, mask = create_placeholder_rectangle(width, height)
                        else:
                            _, mask = create_placeholder_rectangle(width, height)
                        metadata = {"placeholder": True}
                
                child_nodes_data.append({
                    "mask": mask,
                    "bbox": child_bbox,
                    "metadata": metadata,
                    "type": child.type,
                })
            
            # Get container bbox (use node's bbox or parent's)
            container_bbox = (
                node.bbox.get("x", 0),
                node.bbox.get("y", 0),
                node.bbox.get("width", parent_bbox[2]),
                node.bbox.get("height", parent_bbox[3]),
            )
            
            # Optimize children layout
            constraints = node.constraints or {}
            
            # Generate unique save prefix for this node using node path
            # Replace dots and special characters to make valid filename
            save_prefix = node_path.replace(".", "_").replace(" ", "_")
            
            # Extract grandchildren info for proximity ratio calculation
            # For each child, collect its children's bboxes (grandchildren of current container)
            grandchildren_list = []
            for child_result in child_results:
                grandchildren = []
                # Check if child_result has children (from recursive optimization)
                child_children = child_result.get("children", [])
                for grandchild_result in child_children:
                    grandchild_bbox = grandchild_result.get("final_bbox")
                    if grandchild_bbox:
                        if isinstance(grandchild_bbox, (tuple, list)) and len(grandchild_bbox) >= 4:
                            grandchildren.append(tuple(grandchild_bbox[:4]))
                        elif isinstance(grandchild_bbox, dict):
                            grandchildren.append((
                                grandchild_bbox.get("x", 0),
                                grandchild_bbox.get("y", 0),
                                grandchild_bbox.get("width", grandchild_bbox.get("w", 0)),
                                grandchild_bbox.get("height", grandchild_bbox.get("h", 0))
                            ))
                grandchildren_list.append(grandchildren)
            
            config = {
                **self.config.optimization_params,
                "device": self.config.device,
                "debug": self.config.debug,  # Pass debug flag
                "w_similarity": self.config.weights.get("w_similarity", 1.0),
                "w_readability": self.config.weights.get("w_readability", 1.0),
                "w_alignment_consistency": self.config.weights.get("w_alignment_consistency", 1.0),
                "w_alignment_similarity": self.config.weights.get("w_alignment_similarity", params.W_ALIGNMENT_SIMILARITY),
                "w_proximity": self.config.weights.get("w_proximity", params.W_PROXIMITY),
                "container_type": node.type,  # Pass container type (column, row, or layer)
                "grandchildren_list": grandchildren_list,  # Pass grandchildren for proximity calculation
            }
            
            # print(f"[HierarchicalOptimizer] Calling strategy.optimize for container_type={node.type}, num_children={len(child_nodes_data)}")
            # print("child_nodes_data:", child_nodes_data)
            # print("config:", config)
            optimized_bboxes = self.config.strategy.optimize(
                child_nodes_data,
                container_bbox,
                constraints,
                config,
                save_prefix=save_prefix,
            )
            
            # print(f"[HierarchicalOptimizer] Strategy returned optimized_bboxes: {optimized_bboxes}")
            # Calculate actual container bbox based on children's layout results
            # Note: optimized_bboxes are relative to container origin (0,0)
            if optimized_bboxes:
                # Convert bbox to tuple format if needed (handle both dict and tuple formats)
                def bbox_to_tuple(bbox):
                    if isinstance(bbox, dict):
                        return (
                            bbox.get("x", 0),
                            bbox.get("y", 0),
                            bbox.get("width", bbox.get("w", 0)),
                            bbox.get("height", bbox.get("h", 0))
                        )
                    elif isinstance(bbox, (tuple, list)) and len(bbox) >= 4:
                        return tuple(bbox[:4])
                    else:
                        return (0, 0, 0, 0)
                
                bbox_tuples = [bbox_to_tuple(bbox) for bbox in optimized_bboxes]
                
                # print("bbox_tuples:", bbox_tuples)
                # Find the bounding box that contains all children (for position adjustment)
                min_x = min(bbox[0] for bbox in bbox_tuples)
                min_y = min(bbox[1] for bbox in bbox_tuples)
                
                # Adjust children's bboxes: shift them so that min_x and min_y become 0
                # This makes children relative to the new container origin (0,0)
                # For container nodes, preserve their own calculated width/height
                adjusted_bboxes_for_composite = []
                for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)):
                    bbox_tuple = bbox_to_tuple(opt_bbox)
                    # Adjust coordinates: subtract min_x and min_y to start from (0,0)
                    adjusted_x = bbox_tuple[0] - min_x
                    adjusted_y = bbox_tuple[1] - min_y
                    
                    # For container nodes, preserve the width/height calculated from their children
                    # Only update position (x, y), not size (w, h)
                    if child_result.get("children") and child_result.get("final_bbox"):
                        # Child is a container with its own final_bbox calculated from its children
                        child_final = child_result["final_bbox"]
                        if isinstance(child_final, (tuple, list)) and len(child_final) >= 4:
                            # Use the child's own calculated width and height
                            adjusted_bbox = (
                                adjusted_x,
                                adjusted_y,
                                child_final[2],  # Keep child's calculated width
                                child_final[3]   # Keep child's calculated height
                            )
                        else:
                            adjusted_bbox = (adjusted_x, adjusted_y, bbox_tuple[2], bbox_tuple[3])
                    else:
                        # Leaf node or no children: use optimized bbox size
                        adjusted_bbox = (adjusted_x, adjusted_y, bbox_tuple[2], bbox_tuple[3])
                    
                    adjusted_bboxes_for_composite.append(adjusted_bbox)
                    
                    # Store relative coordinates (relative to parent container) in final_bbox
                    # save_result.py will convert to absolute coordinates by adding parent offsets
                    child_result["final_bbox"] = adjusted_bbox
                    if i < len(node.children):
                        node.children[i].final_bbox = adjusted_bbox
                        # print(f"node.children[{i}].final_bbox:", node.children[i].final_bbox)
                
                # Calculate actual container size based on adjusted bboxes
                # This ensures the container size is based on children's preserved widths/heights
                actual_max_x = max(bbox[0] + bbox[2] for bbox in adjusted_bboxes_for_composite)
                actual_max_y = max(bbox[1] + bbox[3] for bbox in adjusted_bboxes_for_composite)
                actual_width = actual_max_x
                actual_height = actual_max_y
                # print(f"actual_container_size: width={actual_width:.2f}, height={actual_height:.2f}")
                
                # Container position: keep original container position
                # Container size: actual size needed to contain all children
                # Children are now adjusted to start from (0,0) relative to container
                actual_container_bbox = (
                    container_bbox[0],  # Keep original x position
                    container_bbox[1],  # Keep original y position
                    actual_width,       # Actual width needed
                    actual_height       # Actual height needed
                )
            else:
                # Fallback to original container bbox if no children
                actual_container_bbox = container_bbox
                adjusted_bboxes_for_composite = optimized_bboxes
                # Update child results with optimized bboxes (no adjustment needed)
                for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)):
                    child_result["final_bbox"] = opt_bbox
                    if i < len(node.children):
                        node.children[i].final_bbox = opt_bbox
            
            # Composite children results (use actual container bbox and adjusted bboxes)
            composite_mask, composite_sdf = self.config.strategy.composite(
                child_nodes_data,
                adjusted_bboxes_for_composite,
                actual_container_bbox,
            )
            
            result["final_bbox"] = actual_container_bbox
            result["composite_mask"] = composite_mask
            result["composite_sdf"] = composite_sdf
            result["children"] = child_results
            
        else:
            # Leaf node: just load data
            handler = self._get_handler(node.type)
            if handler:
                mask, metadata = handler.load(node.__dict__, self.config.base_dir)
                result["mask"] = mask
                result["metadata"] = metadata
                result["final_bbox"] = (
                    node.bbox.get("x", 0),
                    node.bbox.get("y", 0),
                    node.bbox.get("width", 100),
                    node.bbox.get("height", 100),
                )
        
        return result
    
    def _get_handler(self, node_type: str) -> Optional[NodeHandler]:
        """Get handler for node type.
        
        Args:
            node_type: Type of node
        
        Returns:
            Node handler or None
        """
        return self.config.node_handlers.get(node_type)
    
    def _can_use_rule_based_layout(self, node: LayoutNode) -> bool:
        """Check if node can use rule-based layout instead of SDF optimization.
        
        Rule-based layout is much faster and more accurate for simple row/column layouts.
        
        Args:
            node: Layout node to check
        
        Returns:
            True if rule-based layout can be used
        """
        # Check if rule-based is enabled
        if not self.config.use_rule_based:
            return False
        
        # Check if node type is in the allowed list
        if node.type not in self.config.rule_based_types:
            return False
        
        # Need at least 2 children to benefit from rule-based layout
        if not node.children or len(node.children) < 2:
            return False
        
        # Check for complex overlap constraints that require SDF
        constraints = node.constraints or {}
        if 'overlap' in constraints:
            # Overlap constraints need precise collision detection - use SDF
            return False
        
        return True
    
    def _rule_based_layout(self, node: LayoutNode, parent_bbox: Tuple[float, float, float, float],
                          node_path: str) -> Dict[str, Any]:
        """Execute rule-based layout for simple row/column arrangements.
        
        Args:
            node: Layout node to optimize
            parent_bbox: Parent container bounding box (x, y, w, h)
            node_path: Path string for this node
        
        Returns:
            Dictionary with optimization results
        """
        # print(f"[HierarchicalOptimizer] Using rule-based layout for {node.type} node: {node_path}")
        
        result = {
            "type": node.type,
            "bbox": node.bbox,
            "final_bbox": None,
            "image_path": getattr(node, "image_path", None),
        }
        
        # First, recursively optimize all children
        child_results = []
        for i, child in enumerate(node.children):
            child_path = f"{node_path}.child{i}"
            child_result = self._optimize_node(child, parent_bbox, child_path)
            child_results.append(child_result)
        
        # Load child node data
        child_nodes_data = []
        for i, child in enumerate(node.children):
            child_result = child_results[i]
            
            # Get bbox from child result
            if child_result.get("final_bbox") is not None:
                final_bbox = child_result["final_bbox"]
                if isinstance(final_bbox, (tuple, list)) and len(final_bbox) >= 4:
                    child_bbox = {
                        "x": final_bbox[0],
                        "y": final_bbox[1],
                        "width": final_bbox[2],
                        "height": final_bbox[3]
                    }
                else:
                    child_bbox = child.bbox
            else:
                child_bbox = child.bbox
            
            # Get mask
            mask = None
            metadata = {}
            if child_result.get("composite_mask") is not None:
                mask = child_result["composite_mask"]
            else:
                handler = self._get_handler(child.type)
                if handler:
                    mask, metadata = handler.load(child.__dict__, self.config.base_dir)
                else:
                    # Fallback placeholder
                    bbox = child_bbox
                    width = bbox.get("width", 100) if isinstance(bbox, dict) else bbox[2]
                    height = bbox.get("height", 100) if isinstance(bbox, dict) else bbox[3]
                    if child.type in ["layer", "column", "row"]:
                        _, mask = create_placeholder_rounded_rectangle(width, height)
                    else:
                        _, mask = create_placeholder_rectangle(width, height)
                    metadata = {"placeholder": True}
            
            child_nodes_data.append({
                "mask": mask,
                "bbox": child_bbox,
                "metadata": metadata,
                "type": child.type,
            })
        
        # Get container bbox
        container_bbox = (
            node.bbox.get("x", 0),
            node.bbox.get("y", 0),
            node.bbox.get("width", parent_bbox[2]),
            node.bbox.get("height", parent_bbox[3]),
        )
        
        # Use rule-based strategy
        from .strategies import RuleBasedLayoutStrategy
        rule_strategy = RuleBasedLayoutStrategy()
        
        constraints = node.constraints or {}
        config = {
            "container_type": node.type,
        }
        
        # Calculate optimized bboxes using rule-based layout
        optimized_bboxes = rule_strategy.optimize(
            child_nodes_data,
            container_bbox,
            constraints,
            config,
        )
        
        # Adjust bboxes to start from (0, 0) relative to container
        if optimized_bboxes:
            min_x = min(bbox[0] for bbox in optimized_bboxes)
            min_y = min(bbox[1] for bbox in optimized_bboxes)
            
            adjusted_bboxes_for_composite = []
            for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)):
                adjusted_x = opt_bbox[0] - min_x
                adjusted_y = opt_bbox[1] - min_y
                
                # For container nodes, preserve their calculated size
                if child_result.get("children") and child_result.get("final_bbox"):
                    child_final = child_result["final_bbox"]
                    if isinstance(child_final, (tuple, list)) and len(child_final) >= 4:
                        adjusted_bbox = (
                            adjusted_x,
                            adjusted_y,
                            child_final[2],  # Keep child's calculated width
                            child_final[3]   # Keep child's calculated height
                        )
                    else:
                        adjusted_bbox = (adjusted_x, adjusted_y, opt_bbox[2], opt_bbox[3])
                else:
                    adjusted_bbox = (adjusted_x, adjusted_y, opt_bbox[2], opt_bbox[3])
                
                adjusted_bboxes_for_composite.append(adjusted_bbox)
                child_result["final_bbox"] = adjusted_bbox
                if i < len(node.children):
                    node.children[i].final_bbox = adjusted_bbox
            
            # Calculate actual container size
            actual_max_x = max(bbox[0] + bbox[2] for bbox in adjusted_bboxes_for_composite)
            actual_max_y = max(bbox[1] + bbox[3] for bbox in adjusted_bboxes_for_composite)
            actual_width = actual_max_x
            actual_height = actual_max_y
            
            actual_container_bbox = (
                container_bbox[0],
                container_bbox[1],
                actual_width,
                actual_height
            )
        else:
            actual_container_bbox = container_bbox
            adjusted_bboxes_for_composite = optimized_bboxes
            for i, (child_result, opt_bbox) in enumerate(zip(child_results, optimized_bboxes)):
                child_result["final_bbox"] = opt_bbox
                if i < len(node.children):
                    node.children[i].final_bbox = opt_bbox
        
        # Composite children results
        composite_mask, composite_sdf = rule_strategy.composite(
            child_nodes_data,
            adjusted_bboxes_for_composite,
            actual_container_bbox,
        )
        
        result["final_bbox"] = actual_container_bbox
        result["composite_mask"] = composite_mask
        result["composite_sdf"] = composite_sdf
        result["children"] = child_results
        
        return result