"""
布局处理工具模块

提供布局相关处理功能：
- 重叠框检测与去重
- 阅读顺序排序
- OCR Span 与 Layout Block 匹配
- 大面积文本块转换为表格（后处理）

注：底层坐标计算方法（IoU、重叠比例、poly_to_bbox 等）已统一到 coordinate_utils.py
"""
from typing import Dict, List, Any, Tuple, Optional
from loguru import logger
import statistics

# 导入坐标工具（底层坐标计算方法）
try:
    from .coordinate_utils import CoordinateUtils
except ImportError:
    from coordinate_utils import CoordinateUtils


class LayoutUtils:
    """布局处理工具类"""
    
    # ==================== 坐标计算方法（委托给 CoordinateUtils）====================
    
    @staticmethod
    def calculate_iou(bbox1: List[float], bbox2: List[float]) -> float:
        """计算两个 bbox 的 IoU（交并比）- 委托给 CoordinateUtils"""
        return CoordinateUtils.calculate_iou(bbox1, bbox2)
    
    @staticmethod
    def calculate_overlap_ratio(bbox1: List[float], bbox2: List[float]) -> float:
        """计算重叠面积占小框面积的比例 - 委托给 CoordinateUtils"""
        return CoordinateUtils.calculate_overlap_ratio(bbox1, bbox2)
    
    # ==================== 布局处理方法 ====================
    
    @staticmethod
    def remove_overlapping_boxes(
        layout_results: List[Dict[str, Any]],
        iou_threshold: float = 0.8,
        overlap_ratio_threshold: float = 0.8
    ) -> List[Dict[str, Any]]:
        """
        处理重叠的布局框（参考 MinerU 的去重策略）
        
        策略：
        1. 高 IoU 重叠：保留置信度高的框
        2. 包含关系：小框被大框高度包含时，保留大框并扩展边界
        3. 同类型优先合并
        
        Args:
            layout_results: Layout 检测结果列表
            iou_threshold: IoU 阈值，超过此值认为高度重叠
            overlap_ratio_threshold: 重叠面积占小框面积的比例阈值
            
        Returns:
            去重后的布局结果列表
        """
        if not layout_results or len(layout_results) <= 1:
            return layout_results
        
        # 复制列表避免修改原数据
        results = [item.copy() for item in layout_results]
        need_remove = set()
        
        for i in range(len(results)):
            if i in need_remove:
                continue
                
            for j in range(i + 1, len(results)):
                if j in need_remove:
                    continue
                
                bbox1 = results[i].get('bbox', [0, 0, 0, 0])
                bbox2 = results[j].get('bbox', [0, 0, 0, 0])
                
                if len(bbox1) < 4 or len(bbox2) < 4:
                    continue
                
                # 计算 IoU
                iou = LayoutUtils.calculate_iou(bbox1, bbox2)
                
                if iou > iou_threshold:
                    # 高度重叠，保留置信度高的
                    score1 = results[i].get('confidence', results[i].get('score', 0))
                    score2 = results[j].get('confidence', results[j].get('score', 0))
                    
                    if score1 >= score2:
                        need_remove.add(j)
                    else:
                        need_remove.add(i)
                        break  # i 被移除，跳出内层循环
                else:
                    # 检查包含关系
                    overlap_ratio = LayoutUtils.calculate_overlap_ratio(bbox1, bbox2)
                    
                    if overlap_ratio > overlap_ratio_threshold:
                        # 小框被大框高度包含
                        area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
                        area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
                        
                        if area1 <= area2:
                            small_idx, large_idx = i, j
                        else:
                            small_idx, large_idx = j, i
                        
                        # 扩展大框的边界
                        small_bbox = results[small_idx]['bbox']
                        large_bbox = results[large_idx]['bbox']
                        results[large_idx]['bbox'] = [
                            min(small_bbox[0], large_bbox[0]),
                            min(small_bbox[1], large_bbox[1]),
                            max(small_bbox[2], large_bbox[2]),
                            max(small_bbox[3], large_bbox[3])
                        ]
                        need_remove.add(small_idx)
                        
                        if small_idx == i:
                            break  # i 被移除，跳出内层循环
        
        # 返回去重后的结果
        return [results[i] for i in range(len(results)) if i not in need_remove]
    
    @staticmethod
    def convert_large_text_to_table(
        layout_results: List[Dict[str, Any]],
        image_shape: Tuple[int, int],
        min_area_ratio: float = 0.25,
        min_width_ratio: float = 0.4,
        min_height_ratio: float = 0.3
    ) -> List[Dict[str, Any]]:
        """
        将大面积的文本块转换为表格
        
        判断规则：
        1. 面积占比：占页面面积超过 min_area_ratio（默认25%）
        2. 尺寸比例：宽度和高度都超过一定比例（避免细长条）
        3. 不与其他表格重叠：如果已有表格，不转换
        
        Args:
            layout_results: Layout 检测结果列表
            image_shape: 图像尺寸 (height, width)
            min_area_ratio: 最小面积占比（0-1），默认0.25（25%）
            min_width_ratio: 最小宽度占比（0-1），默认0.4（40%）
            min_height_ratio: 最小高度占比（0-1），默认0.3（30%）
            
        Returns:
            转换后的布局结果列表
        """
        if not layout_results:
            return layout_results
        
        img_height, img_width = image_shape
        img_area = img_height * img_width
        
        # 检查是否已有表格
        has_table = any(
            item.get('category', '').lower() in ['table', 'table_body']
            for item in layout_results
        )
        
        # 如果已有表格，不进行转换（避免误判）
        if has_table:
            logger.debug("📋 Page already has table elements, skipping text-to-table conversion")
            return layout_results
        
        # 复制列表避免修改原数据
        results = [item.copy() for item in layout_results]
        converted_count = 0
        
        for item in results:
            category = item.get('category', '').lower()
            
            # 只处理文本类型的元素
            if category not in ['text', 'ocr_text']:
                continue
            
            bbox = item.get('bbox', [0, 0, 0, 0])
            if len(bbox) < 4:
                continue
            
            x1, y1, x2, y2 = bbox[:4]
            width = x2 - x1
            height = y2 - y1
            area = width * height
            
            # 计算占比
            area_ratio = area / img_area if img_area > 0 else 0
            width_ratio = width / img_width if img_width > 0 else 0
            height_ratio = height / img_height if img_height > 0 else 0
            
            # 判断是否满足转换条件
            if (area_ratio >= min_area_ratio and 
                width_ratio >= min_width_ratio and 
                height_ratio >= min_height_ratio):
                
                # 转换为表格
                item['category'] = 'table'
                item['original_category'] = category  # 保留原始类别
                converted_count += 1
                
                logger.info(
                    f"🔄 Converted large text block to table: "
                    f"area={area_ratio:.1%}, size={width_ratio:.1%}×{height_ratio:.1%}, "
                    f"bbox=[{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]"
                )
        
        if converted_count > 0:
            logger.info(f"✅ Converted {converted_count} large text block(s) to table(s)")
        
        return results
    
    @staticmethod
    def sort_elements_by_reading_order(
        elements: List[Dict[str, Any]],
        y_tolerance: float = 15.0
    ) -> List[Dict[str, Any]]:
        """
        根据阅读顺序对元素进行排序，并添加 reading_order 字段
        
        排序规则：
        1. 按Y坐标分行（考虑容差，Y坐标相近的元素视为同一行）
        2. 同一行内按X坐标从左到右排序
        3. 行与行之间按Y坐标从上到下排序
        
        Args:
            elements: 元素列表（坐标已转换为原始图片坐标系）
            y_tolerance: Y坐标容差，在此范围内的元素被视为同一行
            
        Returns:
            排序后的元素列表，每个元素都添加了 reading_order 字段
        """
        if not elements:
            return elements
        
        # 为每个元素提取排序用的坐标
        elements_with_coords = []
        for elem in elements:
            bbox = elem.get('bbox', [0, 0, 0, 0])
            if len(bbox) >= 4:
                y_top = bbox[1]  # 上边界
                x_left = bbox[0]  # 左边界
            else:
                y_top = 0
                x_left = 0
            elements_with_coords.append((elem, y_top, x_left))
        
        # 先按Y坐标排序
        elements_with_coords.sort(key=lambda x: (x[1], x[2]))
        
        # 按Y坐标分行
        rows = []
        current_row = []
        current_row_y = None
        
        for elem, y_top, x_left in elements_with_coords:
            if current_row_y is None:
                # 第一个元素
                current_row.append((elem, x_left))
                current_row_y = y_top
            elif abs(y_top - current_row_y) <= y_tolerance:
                # 同一行
                current_row.append((elem, x_left))
            else:
                # 新的一行
                rows.append(current_row)
                current_row = [(elem, x_left)]
                current_row_y = y_top
        
        # 添加最后一行
        if current_row:
            rows.append(current_row)
        
        # 每行内按X坐标排序，然后展平
        sorted_elements = []
        reading_order = 0
        
        for row in rows:
            # 行内按X坐标排序
            row.sort(key=lambda x: x[1])
            for elem, _ in row:
                # 添加 reading_order 字段
                elem['reading_order'] = reading_order
                sorted_elements.append(elem)
                reading_order += 1
        
        logger.debug(f"📖 Elements sorted by reading order: {len(sorted_elements)} elements")
        return sorted_elements


class SpanMatcher:
    """
    OCR Span 与 Layout Block 匹配器
    
    参考 MinerU 的处理方式：
    1. 整页 OCR 获取所有 spans
    2. 将 spans 匹配到对应的 layout blocks
    
    注：底层坐标计算方法已统一到 CoordinateUtils
    """
    
    # ==================== 坐标计算方法（委托给 CoordinateUtils）====================
    
    @staticmethod
    def calculate_overlap_area_in_bbox1_ratio(
        bbox1: List[float], 
        bbox2: List[float]
    ) -> float:
        """计算 bbox1 被 bbox2 覆盖的面积比例 - 委托给 CoordinateUtils"""
        return CoordinateUtils.calculate_overlap_in_bbox1_ratio(bbox1, bbox2)
    
    @staticmethod
    def poly_to_bbox(poly: List) -> List[float]:
        """将多边形坐标转换为 bbox 格式 - 委托给 CoordinateUtils"""
        return CoordinateUtils.poly_to_bbox(poly)
    
    # ==================== Span 匹配方法 ====================
    
    @staticmethod
    def match_spans_to_blocks(
        ocr_spans: List[Dict[str, Any]],
        layout_blocks: List[Dict[str, Any]],
        overlap_threshold: float = 0.5
    ) -> Dict[int, List[Dict[str, Any]]]:
        """
        将 OCR spans 匹配到 layout blocks
        
        Args:
            ocr_spans: OCR 识别结果列表，每项包含 'bbox'/'poly', 'text', 'confidence'
            layout_blocks: Layout 检测结果列表，每项包含 'bbox', 'category'
            overlap_threshold: 重叠比例阈值，span 被 block 覆盖超过此比例才算匹配
            
        Returns:
            匹配结果字典 {block_index: [matched_spans]}
        """
        matched = {i: [] for i in range(len(layout_blocks))}
        
        for span in ocr_spans:
            # 获取 span 的 bbox
            span_bbox = span.get('bbox', [])
            if not span_bbox:
                continue
            
            # 转换为标准 bbox 格式
            span_bbox = SpanMatcher.poly_to_bbox(span_bbox)
            
            # 找到最佳匹配的 block
            best_match_idx = -1
            best_overlap = 0.0
            
            for block_idx, block in enumerate(layout_blocks):
                block_bbox = block.get('bbox', [0, 0, 0, 0])
                
                overlap = SpanMatcher.calculate_overlap_area_in_bbox1_ratio(
                    span_bbox, block_bbox
                )
                
                if overlap > overlap_threshold and overlap > best_overlap:
                    best_overlap = overlap
                    best_match_idx = block_idx
            
            if best_match_idx >= 0:
                # 创建带绝对坐标的 span 副本
                matched_span = span.copy()
                matched_span['bbox'] = span_bbox  # 确保是标准 bbox 格式
                matched[best_match_idx].append(matched_span)
        
        return matched
    
    @staticmethod
    def merge_spans_to_text(
        spans: List[Dict[str, Any]],
        block_bbox: Optional[List[float]] = None
    ) -> Tuple[str, List[Dict[str, Any]]]:
        """
        将多个 spans 合并为单个文本字符串
        
        参考 MinerU 的 span 合并逻辑：
        1. 按 Y 坐标分行
        2. 同行内按 X 坐标排序
        3. 行间添加换行，词间可能添加空格
        
        Args:
            spans: span 列表
            block_bbox: 所属 block 的 bbox（用于参考）
            
        Returns:
            (merged_text, sorted_spans)
        """
        if not spans:
            return "", []
        
        # 计算 spans 的高度中位数（用于判断同行）
        heights = []
        for span in spans:
            bbox = span.get('bbox', [0, 0, 0, 0])
            if len(bbox) >= 4:
                h = bbox[3] - bbox[1]
                if h > 0:
                    heights.append(h)
        
        if heights:
            median_height = statistics.median(heights)
            y_tolerance = median_height * 0.5
        else:
            y_tolerance = 10
        
        # 为每个 span 添加坐标信息用于排序
        spans_with_coords = []
        for span in spans:
            bbox = span.get('bbox', [0, 0, 0, 0])
            if len(bbox) >= 4:
                y_center = (bbox[1] + bbox[3]) / 2
                x_left = bbox[0]
            else:
                y_center = 0
                x_left = 0
            spans_with_coords.append((span, y_center, x_left))
        
        # 按 Y 坐标分行
        spans_with_coords.sort(key=lambda x: (x[1], x[2]))
        
        lines = []
        current_line = []
        current_line_y = None
        
        for span, y_center, x_left in spans_with_coords:
            if current_line_y is None:
                current_line.append((span, x_left))
                current_line_y = y_center
            elif abs(y_center - current_line_y) <= y_tolerance:
                current_line.append((span, x_left))
            else:
                lines.append(current_line)
                current_line = [(span, x_left)]
                current_line_y = y_center
        
        if current_line:
            lines.append(current_line)
        
        # 合并文本
        text_parts = []
        sorted_spans = []
        
        for line in lines:
            # 行内按 X 坐标排序
            line.sort(key=lambda x: x[1])
            
            line_texts = []
            for span, _ in line:
                text = span.get('text', '')
                if text:
                    line_texts.append(text)
                sorted_spans.append(span)
            
            if line_texts:
                text_parts.append(' '.join(line_texts))
        
        merged_text = '\n'.join(text_parts)
        
        return merged_text, sorted_spans
    
    @staticmethod
    def remove_duplicate_spans(
        spans: List[Dict[str, Any]],
        iou_threshold: float = 0.9
    ) -> List[Dict[str, Any]]:
        """
        移除重复的 spans（高 IoU 重叠）
        
        Args:
            spans: span 列表
            iou_threshold: IoU 阈值
            
        Returns:
            去重后的 spans
        """
        if len(spans) <= 1:
            return spans
        
        result = []
        removed = set()
        
        for i, span1 in enumerate(spans):
            if i in removed:
                continue
            
            bbox1 = span1.get('bbox', [0, 0, 0, 0])
            bbox1 = CoordinateUtils.poly_to_bbox(bbox1)
            
            for j in range(i + 1, len(spans)):
                if j in removed:
                    continue
                
                bbox2 = spans[j].get('bbox', [0, 0, 0, 0])
                bbox2 = CoordinateUtils.poly_to_bbox(bbox2)
                
                iou = CoordinateUtils.calculate_iou(bbox1, bbox2)
                
                if iou > iou_threshold:
                    # 保留置信度高的
                    score1 = span1.get('confidence', span1.get('score', 0))
                    score2 = spans[j].get('confidence', spans[j].get('score', 0))
                    
                    if score1 >= score2:
                        removed.add(j)
                    else:
                        removed.add(i)
                        break
            
            if i not in removed:
                result.append(span1)
        
        return result

