modules/infographics_generator/layout_system/utils/nodes.py

"""
布局系统核心节点定义

基于 template.ebnf 实现的布局系统
支持：
- Flow Layout (ROW/COLUMN)
- Non-Flow Layout (Z-layer)
"""

from abc import ABC, abstractmethod
from typing import List, Optional, Tuple, Any
from dataclasses import dataclass, field
from enum import Enum
import numpy as np


# ============= 枚举定义 =============

class NodeType(Enum):
    """节点类型"""
    GROUP = "GROUP"
    TEXT = "TEXT"
    IMAGE = "IMAGE"
    CHART = "CHART"
    SHAPE = "SHAPE"


class LayoutType(Enum):
    """布局类型"""
    FLOW = "FLOW"
    NON_FLOW = "NON_FLOW"


class FlowDirection(Enum):
    """Flow布局方向"""
    ROW = "ROW"           # 水平排列
    COLUMN = "COLUMN"     # 垂直排列
    # 未来扩展:
    # GRID = "GRID"
    # CIRCULAR = "CIRCULAR"
    # IRREGULAR = "IRREGULAR"


class MainAlignment(Enum):
    """主轴对齐方式"""
    START = "START"
    CENTER = "CENTER"
    END = "END"


class CrossAlignment(Enum):
    """交叉轴对齐方式"""
    START = "START"
    CENTER = "CENTER"
    END = "END"
    STRETCH = "STRETCH"


class PositionAlign(Enum):
    """非流式布局位置对齐"""
    TOP_LEFT = "top-left"
    TOP_CENTER = "top-center"
    TOP_RIGHT = "top-right"
    CENTER_LEFT = "left"
    CENTER = "center"
    CENTER_RIGHT = "right"
    BOTTOM_LEFT = "bottom-left"
    BOTTOM_CENTER = "bottom-center"
    BOTTOM_RIGHT = "bottom-right"


class Alignment(Enum):
    """对齐方式（用于每个节点自己的对齐）"""
    START = "START"
    CENTER = "CENTER"
    END = "END"
    LEFT = "LEFT"
    RIGHT = "RIGHT"
    # 特殊值（用于 layer 子节点）
    BACKGROUND = "BACKGROUND"
    TOP_LEFT = "TOP_LEFT"
    TOP_CENTER = "TOP_CENTER"
    TOP_RIGHT = "TOP_RIGHT"
    BOTTOM_LEFT = "BOTTOM_LEFT"
    BOTTOM_CENTER = "BOTTOM_CENTER"
    BOTTOM_RIGHT = "BOTTOM_RIGHT"


# ============= 数据类定义 =============

@dataclass
class BoundingBox:
    """边界框 (x, y, width, height)"""
    x: float
    y: float
    width: float
    height: float
    
    @property
    def left(self) -> float:
        return self.x
    
    @property
    def right(self) -> float:
        return self.x + self.width
    
    @property
    def top(self) -> float:
        return self.y
    
    @property
    def bottom(self) -> float:
        return self.y + self.height
    
    @property
    def center_x(self) -> float:
        return self.x + self.width / 2
    
    @property
    def center_y(self) -> float:
        return self.y + self.height / 2


@dataclass
class Padding:
    """内边距"""
    top: float = 0
    right: float = 0
    bottom: float = 0
    left: float = 0
    
    @classmethod
    def uniform(cls, value: float) -> 'Padding':
        """创建统一的内边距"""
        return cls(value, value, value, value)
    
    @property
    def horizontal(self) -> float:
        """水平方向总内边距"""
        return self.left + self.right
    
    @property
    def vertical(self) -> float:
        """垂直方向总内边距"""
        return self.top + self.bottom


@dataclass
class FlowAlignment:
    """Flow布局对齐配置"""
    main: MainAlignment = MainAlignment.START
    cross: CrossAlignment = CrossAlignment.START


# ============= 抽象基类 =============

class Node(ABC):
    """节点抽象基类"""
    
    def __init__(self, node_id: str, node_type: NodeType, alignment: Optional[str] = None, parent: Optional['Node'] = None):
        self.id = node_id
        self.type = node_type
        self.alignment = alignment  # 每个节点自己的对齐方式
        self.bbox: Optional[BoundingBox] = None  # 布局计算后的边界框
        self.parent: Optional['Node'] = parent  # 父节点引用
        
    @abstractmethod
    def compute_intrinsic_size(self) -> Tuple[float, float]:
        """
        计算节点的固有尺寸 (width, height)
        
        对于 Leaf Node: 返回内容的实际尺寸
        对于 Non-Leaf Node: 根据子节点和布局规则计算
        """
        pass
    
    @abstractmethod
    def layout(self, x: float, y: float, available_width: Optional[float] = None, 
               available_height: Optional[float] = None) -> BoundingBox:
        """
        执行布局计算
        
        Args:
            x: 起始 x 坐标
            y: 起始 y 坐标
            available_width: 可用宽度（可选）
            available_height: 可用高度（可选）
            
        Returns:
            计算后的边界框
        """
        pass
    
    @abstractmethod
    def to_dict(self) -> dict:
        """序列化为字典"""
        pass


# ============= Leaf Node =============

class LeafNode(Node):
    """叶子节点
    
    叶子节点表示具体的视觉元素，具有固定的尺寸和可选的遮罩
    """
    
    def __init__(self, node_id: str, node_type: NodeType, 
                 width: float, height: float, 
                 mask: Optional[np.ndarray] = None,
                 metadata: Optional[dict] = None,
                 alignment: Optional[str] = None,
                 parent: Optional[Node] = None):
        super().__init__(node_id, node_type, alignment, parent)
        self.width = width
        self.height = height
        self.mask = mask  # 可选的二值遮罩，用于不规则形状
        self.metadata = metadata or {}  # 存储额外的元数据（如content, role, src等）
        
    def compute_intrinsic_size(self) -> Tuple[float, float]:
        """叶子节点的固有尺寸就是其宽高"""
        return (self.width, self.height)
    
    def layout(self, x: float, y: float, available_width: Optional[float] = None,
               available_height: Optional[float] = None) -> BoundingBox:
        """叶子节点的布局很简单，直接放置在指定位置"""
        self.bbox = BoundingBox(x, y, self.width, self.height)
        return self.bbox
    
    def to_dict(self) -> dict:
        result = {
            "id": self.id,
            "type": self.type.value,
            "bbox": {
                "x": self.bbox.x if self.bbox else 0,
                "y": self.bbox.y if self.bbox else 0,
                "width": self.width,
                "height": self.height
            }
        }
        # 添加元数据
        if self.metadata:
            result.update(self.metadata)
        return result


# ============= Non-Leaf Node =============

class GroupNode(Node):
    """组节点（非叶子节点）
    
    组节点包含多个子节点，并根据布局类型对子节点进行排列
    """
    
    def __init__(self, node_id: str, layout_type: LayoutType, 
                 children: List[Node], padding: Optional[Padding] = None,
                 alignment: Optional[str] = None,
                 parent: Optional[Node] = None):
        super().__init__(node_id, NodeType.GROUP, alignment, parent)
        self.layout_type = layout_type
        self.children = children
        # 设置每个子节点的 parent
        for child in self.children:
            child.parent = self
        self.padding = padding or Padding()
        self.mask: Optional[np.ndarray] = None  # 根据子节点mask合并得到
        
        # 布局特定属性（子类设置）
        self.layout_attrs = {}
        
    def compute_intrinsic_size(self) -> Tuple[float, float]:
        """
        根据子节点和布局规则计算固有尺寸
        这个方法会在具体的布局子类中实现
        """
        raise NotImplementedError("Subclass must implement compute_intrinsic_size")
    
    def layout(self, x: float, y: float, available_width: Optional[float] = None,
               available_height: Optional[float] = None) -> BoundingBox:
        """
        执行组节点的布局
        具体的布局算法由子类实现
        """
        raise NotImplementedError("Subclass must implement layout")
    
    def _compute_mask_from_children(self):
        """
        根据所有子节点的mask计算父节点的mask
        
        将每个子节点的mask根据其bbox位置转换到父节点坐标系，
        然后合并所有mask（使用OR操作）
        """
        if not self.bbox or not self.children:
            self.mask = None
            return
        
        # 收集所有有mask的子节点
        children_with_mask = [
            child for child in self.children 
            if child.mask is not None and child.bbox is not None
        ]
        
        if not children_with_mask:
            self.mask = None
            return
        
        # 创建父节点的mask（全零）
        parent_width = int(self.bbox.width)
        parent_height = int(self.bbox.height)
        parent_mask = np.zeros((parent_height, parent_width), dtype=np.uint8)
        
        # 将每个子节点的mask转换到父节点坐标系并合并
        for child in children_with_mask:
            child_mask = child.mask
            child_bbox = child.bbox
            
            # 计算子节点在父节点坐标系中的位置（相对于父节点左上角）
            child_x_in_parent = int(child_bbox.x - self.bbox.x)
            child_y_in_parent = int(child_bbox.y - self.bbox.y)
            child_width = int(child_bbox.width)
            child_height = int(child_bbox.height)
            
            # 确保子节点mask的尺寸与bbox一致
            if child_mask.shape != (child_height, child_width):
                # 如果尺寸不匹配，调整mask尺寸（使用最近邻插值）
                mask_h, mask_w = child_mask.shape
                
                # 计算源坐标
                y_coords = np.clip(
                    (np.arange(child_height) * mask_h / child_height).astype(int),
                    0, mask_h - 1
                )
                x_coords = np.clip(
                    (np.arange(child_width) * mask_w / child_width).astype(int),
                    0, mask_w - 1
                )
                
                # 使用numpy的高级索引进行最近邻插值
                y_indices, x_indices = np.meshgrid(y_coords, x_coords, indexing='ij')
                child_mask = child_mask[y_indices, x_indices]
            
            # 计算在父节点mask中的位置范围
            y_start = max(0, child_y_in_parent)
            y_end = min(parent_height, child_y_in_parent + child_height)
            x_start = max(0, child_x_in_parent)
            x_end = min(parent_width, child_x_in_parent + child_width)
            
            # 计算在子节点mask中的对应范围
            child_y_start = max(0, -child_y_in_parent)
            child_y_end = child_y_start + (y_end - y_start)
            child_x_start = max(0, -child_x_in_parent)
            child_x_end = child_x_start + (x_end - x_start)
            
            # 将子节点mask复制到父节点mask的对应位置（OR操作）
            if (y_end > y_start and x_end > x_start and 
                child_y_end > child_y_start and child_x_end > child_x_start):
                parent_mask[y_start:y_end, x_start:x_end] = np.maximum(
                    parent_mask[y_start:y_end, x_start:x_end],
                    child_mask[child_y_start:child_y_end, child_x_start:child_x_end]
                )
        
        self.mask = parent_mask
    
    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "type": self.type.value,
            "layout": self.layout_type.value,
            "layoutAttrs": self.layout_attrs,
            "children": [child.to_dict() for child in self.children],
            "bbox": {
                "x": self.bbox.x if self.bbox else 0,
                "y": self.bbox.y if self.bbox else 0,
                "width": self.bbox.width if self.bbox else 0,
                "height": self.bbox.height if self.bbox else 0
            }
        }