# Copyright (c) Opendatalab. All rights reserved.
"""
计算表格HTML中每个td单元格的bbox和score
"""
import re
from typing import List, Dict, Optional, Tuple
try:
    from bs4 import BeautifulSoup
except ImportError:
    BeautifulSoup = None
from loguru import logger
import numpy as np


def parse_html_to_grid(html: str) -> List[Dict]:
    """解析HTML表格，返回行列结构"""
    if BeautifulSoup is None:
        logger.warning("BeautifulSoup not available, using simple regex parsing")
        return _parse_html_simple(html)
    
    try:
        soup = BeautifulSoup(html, 'html.parser')
        rows = []
        row_idx = -1
        
        for tr in soup.find_all('tr'):
            row_idx += 1
            cols = []
            col_idx = 0
            
            for cell in tr.find_all(['td', 'th']):
                text = cell.get_text(strip=True)
                rowspan = int(cell.get('rowspan', 1))
                colspan = int(cell.get('colspan', 1))
                
                cols.append({
                    'col_index': col_idx,
                    'rowspan': rowspan,
                    'colspan': colspan,
                    'text': text
                })
                
                col_idx += 1
            
            if cols:
                rows.append({
                    'row_index': row_idx,
                    'cols': cols
                })
        
        return rows
    except Exception as e:
        logger.warning(f"Failed to parse HTML with BeautifulSoup: {e}, using simple parser")
        return _parse_html_simple(html)


def _parse_html_simple(html: str) -> List[Dict]:
    """简单的正则解析作为备选"""
    rows = []
    row_idx = -1
    
    # 提取所有tr块
    tr_pattern = r'<tr>(.*?)</tr>'
    tr_matches = re.findall(tr_pattern, html, re.DOTALL)
    
    for tr_content in tr_matches:
        row_idx += 1
        cols = []
        col_idx = 0
        
        # 提取td/th
        td_pattern = r'<t[dh](?:\s+[^>]*)?>(.*?)</t[dh]>'
        td_matches = re.findall(td_pattern, tr_content, re.DOTALL)
        
        # 提取rowspan和colspan
        td_full_pattern = r'<t[dh]([^>]*)>(.*?)</t[dh]>'
        td_full_matches = re.findall(td_full_pattern, tr_content, re.DOTALL)
        
        for attr_str, content in td_full_matches:
            rowspan = 1
            colspan = 1
            
            rowspan_match = re.search(r'rowspan\s*=\s*(\d+)', attr_str)
            if rowspan_match:
                rowspan = int(rowspan_match.group(1))
            
            colspan_match = re.search(r'colspan\s*=\s*(\d+)', attr_str)
            if colspan_match:
                colspan = int(colspan_match.group(1))
            
            text = re.sub(r'<[^>]+>', '', content).strip()
            
            cols.append({
                'col_index': col_idx,
                'rowspan': rowspan,
                'colspan': colspan,
                'text': text
            })
            
            col_idx += 1
        
        if cols:
            rows.append({
                'row_index': row_idx,
                'cols': cols
            })
    
    return rows


def infer_max_row_col(rows: List[Dict]) -> Tuple[int, int]:
    """推断表格的最大行列数"""
    max_row = len(rows)
    max_col = 0
    
    for row in rows:
        last_col = row['cols'][-1] if row['cols'] else None
        if last_col:
            max_col = max(max_col, last_col['col_index'] + last_col['colspan'])
    
    return max_row, max_col


def uniform_bounds(start: float, end: float, count: int) -> List[float]:
    """均匀分割区间"""
    if count <= 0:
        return [start, end]
    
    step = (end - start) / count
    bounds = [start + i * step for i in range(count + 1)]
    return bounds


def map_ocr_boxes_to_page(ocr_result: List, table_bbox: List, crop_info: Dict = None) -> List[Dict]:
    """将OCR检测框从表格子图坐标映射到页面坐标"""
    if not ocr_result:
        return []
    
    # 提取表格裁剪信息
    # 如果crop_info为None，说明OCR坐标已经是页面坐标（无需转换）
    if crop_info is None:
        crop_xmin = 0
        crop_ymin = 0
    else:
        crop_xmin = crop_info.get('crop_xmin', 0)
        crop_ymin = crop_info.get('crop_ymin', 0)
    
    page_boxes = []
    for item in ocr_result:
        if len(item) < 3:
            continue
        
        dt_box = item[0]
        text = item[1] if len(item) > 1 else ""
        score = item[2] if len(item) > 2 else 0.0
        
        # dt_box可能是多边形或bbox
        if isinstance(dt_box, (list, np.ndarray)):
            box_array = np.array(dt_box)
            # 检查数组形状和大小
            if box_array.ndim == 1 and len(box_array) == 4:  # [x1, y1, x2, y2]
                x1, y1, x2, y2 = float(box_array[0]), float(box_array[1]), float(box_array[2]), float(box_array[3])
            elif box_array.ndim == 2 and box_array.shape[0] == 4 and box_array.shape[1] == 2:  # [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
                x1, y1 = float(box_array[0][0]), float(box_array[0][1])
                x2, y2 = float(box_array[2][0]), float(box_array[2][1])
            else:
                continue
            
            # 映射到页面坐标
            page_x1 = crop_xmin + x1
            page_y1 = crop_ymin + y1
            page_x2 = crop_xmin + x2
            page_y2 = crop_ymin + y2
            
            page_boxes.append({
                'bbox': [page_x1, page_y1, page_x2, page_y2],
                'text': text,
                'score': float(score) if isinstance(score, (int, float)) else 0.0,
                'center_x': (page_x1 + page_x2) / 2,
                'center_y': (page_y1 + page_y2) / 2
            })
    
    return page_boxes


def infer_row_boundaries_from_ocr(page_boxes: List[Dict], max_row: int) -> List[float]:
    """从OCR框推断行边界"""
    if not page_boxes or max_row <= 0:
        return []
    
    # 收集所有y坐标
    y_centers = [box['center_y'] for box in page_boxes]
    if not y_centers:
        return []
    
    y_centers = sorted(y_centers)
    
    # 简单的K-means聚类来分组（或使用更简单的分位数方法）
    if max_row <= 1:
        return [min(y_centers), max(y_centers)]
    
    # 使用分位数方法：将y坐标分成max_row组
    y_sorted = sorted(y_centers)
    step = len(y_sorted) / max_row
    
    boundaries = []
    for i in range(max_row + 1):
        idx = int(i * step)
        if idx >= len(y_sorted):
            idx = len(y_sorted) - 1
        if i == 0:
            boundaries.append(y_sorted[idx] - 5)  # 向上扩展一点
        elif i == max_row:
            boundaries.append(y_sorted[idx] + 5)  # 向下扩展一点
        else:
            boundaries.append(y_sorted[idx])
    
    return sorted(boundaries)


def infer_col_boundaries_from_ocr(page_boxes: List[Dict], max_col: int) -> List[float]:
    """从OCR框推断列边界"""
    if not page_boxes or max_col <= 0:
        return []
    
    x_centers = [box['center_x'] for box in page_boxes]
    if not x_centers:
        return []
    
    x_sorted = sorted(x_centers)
    
    if max_col <= 1:
        return [min(x_sorted), max(x_sorted)]
    
    step = len(x_sorted) / max_col
    
    boundaries = []
    for i in range(max_col + 1):
        idx = int(i * step)
        if idx >= len(x_sorted):
            idx = len(x_sorted) - 1
        if i == 0:
            boundaries.append(x_sorted[idx] - 5)
        elif i == max_col:
            boundaries.append(x_sorted[idx] + 5)
        else:
            boundaries.append(x_sorted[idx])
    
    return sorted(boundaries)


def calculate_iou(bbox1: List[float], bbox2: List[float]) -> float:
    """计算两个bbox的交并比(IoU)"""
    x1_1, y1_1, x2_1, y2_1 = bbox1[:4]
    x1_2, y1_2, x2_2, y2_2 = bbox2[:4]
    
    # 计算交集
    inter_x1 = max(x1_1, x1_2)
    inter_y1 = max(y1_1, y1_2)
    inter_x2 = min(x2_1, x2_2)
    inter_y2 = min(y2_1, y2_2)
    
    if inter_x2 <= inter_x1 or inter_y2 <= inter_y1:
        return 0.0
    
    inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1)
    
    # 计算并集
    area1 = (x2_1 - x1_1) * (y2_1 - y1_1)
    area2 = (x2_2 - x1_2) * (y2_2 - y1_2)
    union_area = area1 + area2 - inter_area
    
    if union_area <= 0:
        return 0.0
    
    return inter_area / union_area


def aggregate_ocr_scores_in_bbox(page_boxes: List[Dict], bbox: List[float], default_score: float = 0.0) -> float:
    """聚合bbox区域内的OCR框的score，使用IoU加权平均"""
    if not page_boxes:
        # 如果没有OCR框，使用提供的default_score（应该是span_score）
        return default_score
    
    # 如果default_score为0或太小，使用更合理的默认值
    if default_score <= 0:
        default_score = 0.8
    
    x1, y1, x2, y2 = bbox[:4]
    weighted_scores = []
    total_weight = 0.0
    
    for box in page_boxes:
        box_x1, box_y1, box_x2, box_y2 = box['bbox'][:4]
        
        # 计算IoU作为权重
        iou = calculate_iou(bbox, [box_x1, box_y1, box_x2, box_y2])
        
        # 只考虑有重叠的框（IoU > 0）
        if iou > 0:
            weight = iou
            box_score = box['score'] if 'score' in box and box['score'] > 0 else default_score
            weighted_scores.append(box_score * weight)
            total_weight += weight
    
    if total_weight > 0 and weighted_scores:
        # 返回加权平均分
        weighted_avg = sum(weighted_scores) / total_weight
        return float(weighted_avg)
    
    return default_score


def compute_table_cells(
    html: str,
    table_bbox: List[float],
    span_score: float = 0.0,
    ocr_result: Optional[List] = None,
    crop_info: Optional[Dict] = None
) -> List[Dict]:
    """
    计算表格HTML中每个td单元格的bbox和score
    
    Args:
        html: 表格HTML字符串
        table_bbox: 表格在页面中的bbox [x1, y1, x2, y2]
        span_score: 表格整体的score
        ocr_result: OCR结果列表（可选），格式为 [[dt_box, text, score], ...]
        crop_info: 裁剪信息（可选），包含crop_xmin, crop_ymin等
    
    Returns:
        单元格列表，每个包含：row_index, col_index, rowspan, colspan, text, bbox, score
    """
    if not html:
        return []
    
    # 安全地检查table_bbox长度
    try:
        bbox_len = len(table_bbox) if hasattr(table_bbox, '__len__') else 0
        if bbox_len < 4:
            return []
    except (TypeError, ValueError):
        return []
    
    # 解析HTML
    rows = parse_html_to_grid(html)
    if not rows:
        return []
    
    max_row, max_col = infer_max_row_col(rows)
    if max_row == 0 or max_col == 0:
        return []
    
    # 安全地提取坐标，处理numpy数组
    try:
        bbox_list = list(table_bbox[:4]) if not isinstance(table_bbox, list) else table_bbox[:4]
        x1, y1, x2, y2 = float(bbox_list[0]), float(bbox_list[1]), float(bbox_list[2]), float(bbox_list[3])
    except (ValueError, TypeError, IndexError) as e:
        logger.warning(f"Error extracting table bbox coordinates: {e}")
        return []
    
    # 计算行边界和列边界
    if ocr_result:
        # 方法B: OCR引导的分割
        page_boxes = map_ocr_boxes_to_page(ocr_result, table_bbox, crop_info)
        
        if page_boxes:
            row_bounds = infer_row_boundaries_from_ocr(page_boxes, max_row)
            col_bounds = infer_col_boundaries_from_ocr(page_boxes, max_col)
        else:
            # 回退到均匀分割
            row_bounds = uniform_bounds(y1, y2, max_row)
            col_bounds = uniform_bounds(x1, x2, max_col)
    else:
        # 方法A: 均匀分割
        row_bounds = uniform_bounds(y1, y2, max_row)
        col_bounds = uniform_bounds(x1, x2, max_col)
    
    # 如果边界数量不足，补充
    while len(row_bounds) < max_row + 1:
        if len(row_bounds) == 0:
            row_bounds = [y1, y2]
        else:
            step = (y2 - y1) / max_row
            row_bounds = [y1 + i * step for i in range(max_row + 1)]
        break
    
    while len(col_bounds) < max_col + 1:
        if len(col_bounds) == 0:
            col_bounds = [x1, x2]
        else:
            step = (x2 - x1) / max_col
            col_bounds = [x1 + i * step for i in range(max_col + 1)]
        break
    
    # 生成单元格bbox和score
    cells = []
    page_boxes_cache = map_ocr_boxes_to_page(ocr_result, table_bbox, crop_info) if ocr_result else []
    
    for row in rows:
        for col in row['cols']:
            row_idx = row['row_index']
            col_idx = col['col_index']
            
            # 计算bbox（保留为整数）
            try:
                row_start = row_idx
                row_end = min(row_idx + col['rowspan'], len(row_bounds) - 1)
                col_start = col_idx
                col_end = min(col_idx + col['colspan'], len(col_bounds) - 1)
                
                cell_y1 = row_bounds[row_start] if row_start < len(row_bounds) else y1
                cell_y2 = row_bounds[row_end] if row_end < len(row_bounds) else y2
                cell_x1 = col_bounds[col_start] if col_start < len(col_bounds) else x1
                cell_x2 = col_bounds[col_end] if col_end < len(col_bounds) else x2
                
                # bbox坐标四舍五入为整数
                cell_bbox = [
                    int(round(cell_x1)), 
                    int(round(cell_y1)), 
                    int(round(cell_x2)), 
                    int(round(cell_y2))
                ]
            except (IndexError, ValueError) as e:
                logger.warning(f"Error calculating cell bbox: {e}, using table bbox")
                cell_bbox = [
                    int(round(x1)), 
                    int(round(y1)), 
                    int(round(x2)), 
                    int(round(y2))
                ]
            
            # ============================================
            # 单元格score计算逻辑
            # ============================================
            # 基于spans['score']的计算公式调整cell['score']
            # 
            # 公式说明：
            # cell_score = base_score * content_factor * structure_factor * quality_factor
            #
            # 其中：
            # base_score: 基于spans['score']的基础置信度（通过IoU加权OCR或直接继承）
            # content_factor: 内容因子（根据文本内容调整）
            # structure_factor: 结构因子（根据单元格结构合理性调整）
            # quality_factor: 质量因子（根据单元格bbox质量调整）
            # ============================================
            
            # 先使用临时浮点bbox计算score（用于IoU计算）
            temp_bbox = [float(cell_bbox[0]), float(cell_bbox[1]), float(cell_bbox[2]), float(cell_bbox[3])]
            
            # 获取单元格文本
            cell_text = col['text'].strip()
            
            # -----------------------------------------------------
            # 步骤1: 计算base_score（基础置信度）
            # -----------------------------------------------------
            # 基于spans['score']的逻辑：spans['score'] = layout_det['score']
            # 对于cells，我们使用表格整体的span_score作为基准
            # 然后通过OCR结果（如果有）进行细粒度调整
            # -----------------------------------------------------
            # 处理span_score为0或未定义的情况，使用合理的默认值
            if span_score <= 0:
                logger.warning(f"Invalid span_score: {span_score}, using default 0.8")
                span_score = 0.8  # 使用合理的默认值
            
            if page_boxes_cache:
                # 如果有OCR结果，使用IoU加权平均（与方法中spans的处理逻辑一致）
                base_score = aggregate_ocr_scores_in_bbox(page_boxes_cache, temp_bbox, span_score)
            else:
                # 如果没有OCR结果，直接使用span_score（类似于spans的处理）
                base_score = span_score
            
            # -----------------------------------------------------
            # 步骤2: 计算content_factor（内容因子）
            # -----------------------------------------------------
            # 根据单元格是否有文本内容调整置信度
            # 有文本 -> 置信度更高，空单元格 -> 置信度较低
            # -----------------------------------------------------
            if cell_text:
                # 有文本：保持或稍微提升置信度
                # 文本越短，可能是标题或标签，置信度稍低
                # 文本较长，可能是内容单元格，置信度稍高
                text_length = len(cell_text)
                if text_length < 3:
                    content_factor = 0.9  # 超短文本（可能是编号、符号等）
                elif text_length < 10:
                    content_factor = 1.0  # 短文本（标题、标签等）
                else:
                    content_factor = 1.05  # 长文本（内容单元格）
                # 限制content_factor在合理范围
                content_factor = min(1.1, content_factor)
            else:
                # 空单元格：显著降低置信度（但不为0，因为空单元格也可能是合理的）
                content_factor = 0.35
            
            # -----------------------------------------------------
            # 步骤3: 计算structure_factor（结构因子）
            # -----------------------------------------------------
            # 根据单元格的跨行跨列情况调整置信度
            # 跨行列的单元格通常结构更复杂，可能需要不同的置信度
            # -----------------------------------------------------
            if col['rowspan'] > 1 or col['colspan'] > 1:
                # 跨行列单元格：保持正常置信度
                structure_factor = 1.0
            else:
                # 普通单元格：保持正常置信度
                structure_factor = 1.0
            
            # -----------------------------------------------------
            # 步骤4: 计算quality_factor（质量因子）
            # -----------------------------------------------------
            # 根据单元格bbox大小和质量调整置信度
            # 太小的单元格可能是噪声，置信度降低
            # -----------------------------------------------------
            cell_width = cell_bbox[2] - cell_bbox[0]
            cell_height = cell_bbox[3] - cell_bbox[1]
            cell_area = cell_width * cell_height
            
            # 单元格bbox质量评估
            if cell_area < 100:  # 面积小于100像素 -> 噪声可能性高
                quality_factor = 0.5
            elif cell_area < 400:  # 面积在100-400像素 -> 可疑
                quality_factor = 0.75
            else:  # 面积>=400像素 -> 正常
                quality_factor = 1.0
            
            # -----------------------------------------------------
            # 步骤5: 综合计算最终cell_score
            # -----------------------------------------------------
            # 应用公式：cell_score = base_score * content_factor * structure_factor * quality_factor
            # 然后限制在[0.0, 1.0]范围内
            # -----------------------------------------------------
            cell_score = base_score * content_factor * structure_factor * quality_factor
            
            # 确保score在合理范围内 [0.0, 1.0]
            cell_score = max(0.0, min(1.0, float(cell_score)))
            
            # 调试日志（第一个单元格）
            if row_idx == 0 and col_idx == 0:
                logger.debug(f"Cell score calculation: base={base_score:.4f}, "
                           f"content={content_factor:.4f}, quality={quality_factor:.4f}, "
                           f"final={cell_score:.4f}, text='{cell_text[:20]}'")
            
            # -----------------------------------------------------
            # 计算公式总结
            # -----------------------------------------------------
            # cell_score = base_score * content_factor * structure_factor * quality_factor
            #
            # 参数说明：
            # - base_score: 
            #     * 有OCR: IoU加权的OCR score（继承spans逻辑）
            #     * 无OCR: span_score（直接继承）
            # - content_factor:
            #     * 空单元格: 0.35
            #     * 超短文本(<3): 0.9
            #     * 短文本(3-10): 1.0
            #     * 长文本(>=10): 1.05 (上限1.1)
            # - structure_factor: 1.0 (当前统一，可根据需要调整)
            # - quality_factor:
            #     * area < 100: 0.5
            #     * 100 <= area < 400: 0.75
            #     * area >= 400: 1.0
            # -----------------------------------------------------
            
            cells.append({
                'row_index': row_idx,
                'col_index': col_idx,
                'rowspan': col['rowspan'],
                'colspan': col['colspan'],
                'text': col['text'],
                'bbox': cell_bbox,  # 已经是整数列表
                'score': round(cell_score, 4)  # 保留4位小数
            })
    
    return cells