src/layout_engine.py

"""
Layout engine for positioning and arranging infographic elements
"""
from typing import Dict, List, Tuple, Optional
import math
import logging
from dataclasses import dataclass

logger = logging.getLogger(__name__)

@dataclass
class LayoutElement:
    """Represents a positioned element in the layout"""
    id: str
    type: str
    content: str
    x: int
    y: int
    width: int
    height: int
    priority: int
    styling: Dict

@dataclass
class LayoutGrid:
    """Grid system for organizing layout"""
    columns: int
    rows: int
    cell_width: int
    cell_height: int
    gap: int

class LayoutEngine:
    """Engine for creating and managing infographic layouts"""
    
    def __init__(self):
        """Initialize layout engine"""
        self.current_layout = None
        self.elements = []
        logger.info("Layout engine initialized")
    
    def create_layout(self, styled_content: Dict) -> Dict:
        """
        Create complete layout from styled content
        
        Args:
            styled_content: Content with applied template styling
            
        Returns:
            Complete layout specification
        """
        layout_type = styled_content.get('layout_type', 'Vertical')
        design_specs = styled_content.get('design_specs', {})
        
        # Initialize layout canvas
        canvas_size = design_specs.get('canvas_size', (1080, 1920))
        margins = design_specs.get('margins', {'top': 60, 'bottom': 60, 'left': 60, 'right': 60})
        
        layout = {
            'type': layout_type,
            'canvas_width': canvas_size[0],
            'canvas_height': canvas_size[1],
            'content_area': {
                'x': margins['left'],
                'y': margins['top'],
                'width': canvas_size[0] - margins['left'] - margins['right'],
                'height': canvas_size[1] - margins['top'] - margins['bottom']
            },
            'elements': [],
            'grid': self._create_grid_system(layout_type, canvas_size, margins, design_specs),
            'flow': self._create_flow_system(layout_type)
        }
        
        # Position elements based on layout type
        if layout_type == 'Vertical':
            layout['elements'] = self._create_vertical_layout(styled_content, layout)
        elif layout_type == 'Horizontal':
            layout['elements'] = self._create_horizontal_layout(styled_content, layout)
        elif layout_type == 'Grid':
            layout['elements'] = self._create_grid_layout(styled_content, layout)
        elif layout_type == 'Flow':
            layout['elements'] = self._create_flow_layout(styled_content, layout)
        else:
            layout['elements'] = self._create_vertical_layout(styled_content, layout)
        
        self.current_layout = layout
        return layout
    
    def _create_grid_system(self, layout_type: str, canvas_size: Tuple[int, int], 
                           margins: Dict, design_specs: Dict) -> LayoutGrid:
        """Create grid system for layout"""
        grid_specs = design_specs.get('grid', {'columns': 1, 'rows': 'auto', 'gap': 30})
        
        content_width = canvas_size[0] - margins['left'] - margins['right']
        content_height = canvas_size[1] - margins['top'] - margins['bottom']
        
        columns = grid_specs['columns']
        gap = grid_specs['gap']
        
        cell_width = (content_width - (columns - 1) * gap) // columns
        cell_height = 200  # Default cell height, can be adjusted
        
        return LayoutGrid(
            columns=columns,
            rows=grid_specs.get('rows', 'auto'),
            cell_width=cell_width,
            cell_height=cell_height,
            gap=gap
        )
    
    def _create_flow_system(self, layout_type: str) -> Dict:
        """Create flow system for dynamic positioning"""
        return {
            'direction': 'vertical' if layout_type in ['Vertical', 'Flow'] else 'horizontal',
            'wrap': layout_type == 'Flow',
            'spacing': 'auto',
            'alignment': 'start'
        }
    
    def _create_vertical_layout(self, styled_content: Dict, layout: Dict) -> List[LayoutElement]:
        """Create vertical layout arrangement"""
        elements = []
        current_y = layout['content_area']['y']
        content_width = layout['content_area']['width']
        content_x = layout['content_area']['x']
        
        # Add title
        title = styled_content.get('title', {})
        if title.get('text'):
            title_element = LayoutElement(
                id='title',
                type='title',
                content=title['text'],
                x=content_x,
                y=current_y,
                width=content_width,
                height=self._calculate_text_height(title['text'], title.get('font', ('Arial', 32, 'bold')), content_width),
                priority=10,
                styling=title
            )
            elements.append(title_element)
            current_y += title_element.height + title.get('margin', 40)
        
        # Add sections
        sections = styled_content.get('sections', [])
        for i, section in enumerate(sections):
            section_height = self._calculate_section_height(section, content_width)
            
            section_element = LayoutElement(
                id=f'section_{section.get("id", i)}',
                type='section',
                content=section.get('condensed_text', section.get('content', '')),
                x=content_x,
                y=current_y,
                width=content_width,
                height=section_height,
                priority=section.get('priority', 5),
                styling=section.get('styling', {})
            )
            elements.append(section_element)
            current_y += section_height + section.get('styling', {}).get('margin', 30)
        
        # Add visual elements
        visual_elements = styled_content.get('visual_elements', [])
        for element in visual_elements:
            if element.get('placement') == 'body':
                visual_element = self._create_visual_element(element, content_x, current_y, content_width)
                elements.append(visual_element)
                current_y += visual_element.height + 20
        
        return elements
    
    def _create_horizontal_layout(self, styled_content: Dict, layout: Dict) -> List[LayoutElement]:
        """Create horizontal layout arrangement"""
        elements = []
        content_area = layout['content_area']
        
        # Title spans full width
        title = styled_content.get('title', {})
        current_y = content_area['y']
        
        if title.get('text'):
            title_element = LayoutElement(
                id='title',
                type='title',
                content=title['text'],
                x=content_area['x'],
                y=current_y,
                width=content_area['width'],
                height=self._calculate_text_height(title['text'], title.get('font', ('Arial', 32, 'bold')), content_area['width']),
                priority=10,
                styling=title
            )
            elements.append(title_element)
            current_y += title_element.height + title.get('margin', 40)
        
        # Split remaining sections into two columns
        sections = styled_content.get('sections', [])
        column_width = (content_area['width'] - 40) // 2  # 40px gap between columns
        
        left_column_y = current_y
        right_column_y = current_y
        
        for i, section in enumerate(sections):
            section_height = self._calculate_section_height(section, column_width)
            
            if i % 2 == 0:  # Left column
                x = content_area['x']
                y = left_column_y
                left_column_y += section_height + 30
            else:  # Right column
                x = content_area['x'] + column_width + 40
                y = right_column_y
                right_column_y += section_height + 30
            
            section_element = LayoutElement(
                id=f'section_{section.get("id", i)}',
                type='section',
                content=section.get('condensed_text', section.get('content', '')),
                x=x,
                y=y,
                width=column_width,
                height=section_height,
                priority=section.get('priority', 5),
                styling=section.get('styling', {})
            )
            elements.append(section_element)
        
        return elements
    
    def _create_grid_layout(self, styled_content: Dict, layout: Dict) -> List[LayoutElement]:
        """Create grid layout arrangement"""
        elements = []
        grid = layout['grid']
        content_area = layout['content_area']
        
        # Title at top
        title = styled_content.get('title', {})
        current_y = content_area['y']
        
        if title.get('text'):
            title_element = LayoutElement(
                id='title',
                type='title',
                content=title['text'],
                x=content_area['x'],
                y=current_y,
                width=content_area['width'],
                height=80,
                priority=10,
                styling=title
            )
            elements.append(title_element)
            current_y += 120
        
        # Arrange sections in grid
        sections = styled_content.get('sections', [])
        grid_start_y = current_y
        
        for i, section in enumerate(sections):
            row = i // grid.columns
            col = i % grid.columns
            
            x = content_area['x'] + col * (grid.cell_width + grid.gap)
            y = grid_start_y + row * (grid.cell_height + grid.gap)
            
            section_element = LayoutElement(
                id=f'section_{section.get("id", i)}',
                type='section',
                content=section.get('condensed_text', section.get('content', '')),
                x=x,
                y=y,
                width=grid.cell_width,
                height=grid.cell_height,
                priority=section.get('priority', 5),
                styling=section.get('styling', {})
            )
            elements.append(section_element)
        
        return elements
    
    def _create_flow_layout(self, styled_content: Dict, layout: Dict) -> List[LayoutElement]:
        """Create flowing layout arrangement"""
        elements = []
        content_area = layout['content_area']
        
        # Title
        title = styled_content.get('title', {})
        current_y = content_area['y']
        
        if title.get('text'):
            title_element = LayoutElement(
                id='title',
                type='title',
                content=title['text'],
                x=content_area['x'],
                y=current_y,
                width=content_area['width'],
                height=80,
                priority=10,
                styling=title
            )
            elements.append(title_element)
            current_y += 100
        
        # Flow sections dynamically
        sections = styled_content.get('sections', [])
        current_x = content_area['x']
        row_height = 0
        max_width = content_area['width']
        
        for i, section in enumerate(sections):
            section_width = min(400, max_width // 2)  # Adaptive width
            section_height = self._calculate_section_height(section, section_width)
            
            # Check if we need to wrap to next row
            if current_x + section_width > content_area['x'] + max_width:
                current_x = content_area['x']
                current_y += row_height + 30
                row_height = 0
            
            section_element = LayoutElement(
                id=f'section_{section.get("id", i)}',
                type='section',
                content=section.get('condensed_text', section.get('content', '')),
                x=current_x,
                y=current_y,
                width=section_width,
                height=section_height,
                priority=section.get('priority', 5),
                styling=section.get('styling', {})
            )
            elements.append(section_element)
            
            current_x += section_width + 20
            row_height = max(row_height, section_height)
        
        return elements
    
    def _calculate_text_height(self, text: str, font: Tuple[str, int, str], width: int) -> int:
        """Calculate approximate text height"""
        if not text:
            return 0
        
        font_size = font[1] if len(font) > 1 else 16
        chars_per_line = width // (font_size * 0.6)  # Rough approximation
        lines = max(1, len(text) / chars_per_line)
        line_height = font_size * 1.4  # Standard line height
        
        return int(lines * line_height)
    
    def _calculate_section_height(self, section: Dict, width: int) -> int:
        """Calculate section height based on content"""
        content = section.get('condensed_text', section.get('content', ''))
        styling = section.get('styling', {})
        font = styling.get('font', ('Arial', 16, 'normal'))
        
        base_height = self._calculate_text_height(content, font, width)
        padding = styling.get('padding', 20)
        
        return base_height + padding * 2
    
    def _create_visual_element(self, element: Dict, x: int, y: int, max_width: int) -> LayoutElement:
        """Create visual element layout"""
        element_type = element.get('type', 'icon')
        importance = element.get('importance', 5)
        
        # Size based on importance
        if element_type == 'chart':
            width = min(max_width, 300)
            height = 200
        elif element_type == 'icon':
            size = 32 + (importance * 4)
            width = height = size
        else:
            width = min(max_width, 250)
            height = 150
        
        return LayoutElement(
            id=f'visual_{element.get("type", "element")}_{id(element)}',
            type=element_type,
            content=element.get('description', ''),
            x=x,
            y=y,
            width=width,
            height=height,
            priority=importance,
            styling=element.get('styling', {})
        )
    
    def optimize_layout(self, layout: Dict) -> Dict:
        """Optimize layout for better visual balance"""
        elements = layout.get('elements', [])
        
        # Remove overlapping elements
        elements = self._resolve_overlaps(elements)
        
        # Balance visual weight
        elements = self._balance_visual_weight(elements, layout)
        
        # Ensure minimum spacing
        elements = self._enforce_minimum_spacing(elements)
        
        layout['elements'] = elements
        return layout
    
    def _resolve_overlaps(self, elements: List[LayoutElement]) -> List[LayoutElement]:
        """Resolve overlapping elements"""
        for i, elem1 in enumerate(elements):
            for j, elem2 in enumerate(elements[i+1:], i+1):
                if self._elements_overlap(elem1, elem2):
                    # Move the lower priority element
                    if elem1.priority < elem2.priority:
                        elem1.y = elem2.y + elem2.height + 20
                    else:
                        elem2.y = elem1.y + elem1.height + 20
        
        return elements
    
    def _elements_overlap(self, elem1: LayoutElement, elem2: LayoutElement) -> bool:
        """Check if two elements overlap"""
        return not (elem1.x + elem1.width <= elem2.x or 
                   elem2.x + elem2.width <= elem1.x or 
                   elem1.y + elem1.height <= elem2.y or 
                   elem2.y + elem2.height <= elem1.y)
    
    def _balance_visual_weight(self, elements: List[LayoutElement], layout: Dict) -> List[LayoutElement]:
        """Balance visual weight of elements"""
        # Sort by priority
        elements.sort(key=lambda x: x.priority, reverse=True)
        
        # Adjust positions for better balance
        canvas_center_x = layout['canvas_width'] // 2
        
        for element in elements:
            if element.type == 'title':
                # Center titles
                element.x = canvas_center_x - element.width // 2
        
        return elements
    
    def _enforce_minimum_spacing(self, elements: List[LayoutElement]) -> List[LayoutElement]:
        """Ensure minimum spacing between elements"""
        min_spacing = 15
        
        # Sort by y position
        elements.sort(key=lambda x: x.y)
        
        for i in range(len(elements) - 1):
            current = elements[i]
            next_elem = elements[i + 1]
            
            required_y = current.y + current.height + min_spacing
            if next_elem.y < required_y:
                next_elem.y = required_y
        
        return elements
    
    def get_layout_bounds(self, layout: Dict) -> Dict:
        """Get the bounds of the entire layout"""
        elements = layout.get('elements', [])
        
        if not elements:
            return {'x': 0, 'y': 0, 'width': 0, 'height': 0}
        
        min_x = min(elem.x for elem in elements)
        min_y = min(elem.y for elem in elements)
        max_x = max(elem.x + elem.width for elem in elements)
        max_y = max(elem.y + elem.height for elem in elements)
        
        return {
            'x': min_x,
            'y': min_y,
            'width': max_x - min_x,
            'height': max_y - min_y
        }
    
    def export_layout_data(self, layout: Dict) -> Dict:
        """Export layout data for image generation"""
        return {
            'canvas': {
                'width': layout['canvas_width'],
                'height': layout['canvas_height'],
                'background': '#ffffff'
            },
            'elements': [
                {
                    'id': elem.id,
                    'type': elem.type,
                    'content': elem.content,
                    'position': {'x': elem.x, 'y': elem.y},
                    'size': {'width': elem.width, 'height': elem.height},
                    'priority': elem.priority,
                    'styling': elem.styling
                }
                for elem in layout.get('elements', [])
            ],
            'bounds': self.get_layout_bounds(layout)
        }