import cv2
import numpy as np
import os
import re
import glob

def clean_filename(text):
    """清理文件名"""
    text = text.lower().strip()
    text = re.sub(r'[\s\W]+', '_', text)
    return text.strip('_')

def sort_contours(cnts, method="left-to-right"):
    """
    对轮廓进行排序。
    对于 Grid 布局，我们需要 'top-to-bottom' 然后 'left-to-right' 的混合排序。
    """
    if not cnts:
        return [], []

    # 获取每个轮廓的 Bounding Box
    boundingBoxes = [cv2.boundingRect(c) for c in cnts]
    
    # 将轮廓和bbox打包
    cnts_boxes = list(zip(cnts, boundingBoxes))
    
    # 1. 按照 Y 坐标（从上到下）进行初步排序
    # key: y
    cnts_boxes.sort(key=lambda b: b[1][1])
    
    # 2. 分行处理
    # 由于手工画线或扫描误差，同一行的y坐标可能不完全相同。
    # 我们需要设定一个阈值，认为y坐标相近的是“同一行”。
    rows = []
    current_row = []
    if cnts_boxes:
        # 以第一个轮廓的高度作为参考阈值
        ref_h = cnts_boxes[0][1][3]
        tolerance = ref_h * 0.5 # 容差设为高度的一半
        
        last_y = cnts_boxes[0][1][1]
        
        for c, box in cnts_boxes:
            y = box[1]
            if y <= last_y + tolerance:
                current_row.append((c, box))
            else:
                # 新的一行
                rows.append(current_row)
                current_row = [(c, box)]
                last_y = y
        # 添加最后一行
        if current_row:
            rows.append(current_row)

    # 3. 对每一行内部，按照 X 坐标（从左到右）排序
    final_sorted = []
    row_counts = []
    for i, row in enumerate(rows):
        # key: x
        row.sort(key=lambda b: b[1][0])
        row_counts.append(len(row))
        for item in row:
            final_sorted.append(item[1]) # 只返回 bbox (x, y, w, h)

    return final_sorted, row_counts

def uniform_grid_split(img, expected_cols=6, expected_rows=4, margin_percent=0.02):
    """
    均匀分割方法：直接按照预期的行列数均匀分割图像
    适用于网格线不连续或没有明显网格线的情况
    
    Args:
        img: 输入图像
        expected_cols: 期望的列数
        expected_rows: 期望的行数
        margin_percent: 边缘裁剪比例（去除可能的边框）
    
    Returns:
        排序好的 (x, y, w, h) 列表
    """
    h_img, w_img = img.shape[:2]
    
    # 去除边缘
    margin_x = int(w_img * margin_percent)
    margin_y = int(h_img * margin_percent)
    
    effective_width = w_img - 2 * margin_x
    effective_height = h_img - 2 * margin_y
    
    # 计算每个单元格的尺寸
    cell_width = effective_width // expected_cols
    cell_height = effective_height // expected_rows
    
    boxes = []
    for row in range(expected_rows):
        for col in range(expected_cols):
            x = margin_x + col * cell_width
            y = margin_y + row * cell_height
            boxes.append((x, y, cell_width, cell_height))
    
    return boxes

def detect_grid_cells_with_lines(img, expected_cols=6, expected_rows=4):
    """
    通过形态学操作检测网格线，并提取每个格子的坐标
    """
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 二值化 (反转：背景黑，内容/线白)
    # 使用自适应阈值来应对光照或颜色不均
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                   cv2.THRESH_BINARY_INV, 11, 2)

    # 定义结构元素 (Kernel) - 增大kernel以更好地检测断裂的线
    h_img, w_img = img.shape[:2]
    # 水平线 Kernel: 宽度长，高度为1
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (w_img // 15, 1))
    # 垂直线 Kernel: 宽度为1，高度长
    vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, h_img // 15))

    # 1. 提取水平线
    detect_horizontal = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
    
    # 2. 提取垂直线
    detect_vertical = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, vertical_kernel, iterations=2)

    # 3. 合并网格线
    grid_mask = cv2.addWeighted(detect_horizontal, 0.5, detect_vertical, 0.5, 0)
    _, grid_mask = cv2.threshold(grid_mask, 0, 255, cv2.THRESH_BINARY)
    
    # 更强的膨胀操作，连接断裂的网格线
    kernel_dilate = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    grid_mask = cv2.dilate(grid_mask, kernel_dilate, iterations=3)
    
    # 闭运算，进一步连接断裂
    kernel_close = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 7))
    grid_mask = cv2.morphologyEx(grid_mask, cv2.MORPH_CLOSE, kernel_close, iterations=2)

    # 4. 寻找所有的“洞”（即单元格）
    # 我们通过查找 grid_mask 的轮廓，通常很难直接找到内部的矩形。
    # 更好的方法是：找出网格线轮廓，画在全黑背景上，然后寻找连通组件，或者反转图片找白色方块。
    
    # 这里我们采用“反转 mask”法：网格线是黑，格子是白
    contours_mask = cv2.bitwise_not(grid_mask)
    contours, _ = cv2.findContours(contours_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 计算期望的单元格面积
    expected_cell_area = (w_img * h_img) / (expected_cols * expected_rows)
    
    # 过滤微小的噪点轮廓，同时也过滤过大的轮廓
    # 放宽过滤条件以捕获更多单元格
    min_area = expected_cell_area / 20  # 单元格面积的 1/20 (之前是1/10)
    max_area = expected_cell_area * 3   # 单元格面积的3倍 (之前是2倍)
    
    # 调试信息
    print(f"    [调试] 图像尺寸: {w_img}x{h_img}, 检测到轮廓: {len(contours)}个")
    print(f"    [调试] 期望单元格面积: {expected_cell_area:.0f}, 过滤范围: {min_area:.0f}-{max_area:.0f}")
    
    # 统计被过滤掉的轮廓
    filtered_out = []
    valid_contours = []
    for c in contours:
        area = cv2.contourArea(c)
        if min_area < area < max_area:
            valid_contours.append(c)
        else:
            x, y, w, h = cv2.boundingRect(c)
            filtered_out.append((area, x, y, w, h))
    
    if filtered_out:
        print(f"    [调试] 被过滤掉 {len(filtered_out)} 个轮廓:")
        for area, x, y, w, h in sorted(filtered_out, key=lambda t: t[0], reverse=True)[:5]:
            print(f"      - 面积={area:.0f}, 位置=({x},{y}), 尺寸={w}x{h}")

    # 排序：确保顺序是 左->右，上->下
    sorted_boxes, row_counts = sort_contours(valid_contours)
    
    return sorted_boxes, row_counts

def detect_grid_cells(img, expected_cols=6, expected_rows=4):
    """
    鲁棒的网格检测方法：首先尝试检测网格线，如果失败则使用均匀分割
    """
    expected_count = expected_cols * expected_rows
    
    # 方法1: 尝试检测网格线
    sorted_boxes, row_counts = detect_grid_cells_with_lines(img, expected_cols, expected_rows)
    
    # 严格检查：必须恰好检测到期望数量的单元格
    if len(sorted_boxes) != expected_count:
        print(f"  网格线检测不理想（检测到 {len(sorted_boxes)} 个单元格，期望 {expected_count} 个）")
        if row_counts:
            row_info = ", ".join([f"第{i+1}行: {count}个" for i, count in enumerate(row_counts)])
            print(f"  检测到的行分布: {row_info}")
        print(f"  切换到均匀分割模式...")
        sorted_boxes = uniform_grid_split(img, expected_cols, expected_rows)
        # 均匀分割时，打印每行的单元格数
        print(f"  均匀分割结果：每行 {expected_cols} 个单元格，共 {expected_rows} 行")
    else:
        print(f"  ✓ 成功检测到 {len(sorted_boxes)} 个网格单元格（符合预期）")
        # 打印每行的单元格数量
        if row_counts:
            row_info = ", ".join([f"第{i+1}行: {count}个" for i, count in enumerate(row_counts)])
            print(f"  行分布: {row_info}")
    
    return sorted_boxes

def is_likely_text_region(img_region, thresh_region):
    """
    判断一个区域是否可能是文字
    文字的特征：
    1. 主要是黑色或深色
    2. 高度较小
    3. 像素密度适中（不是纯色块）
    """
    if img_region.shape[0] == 0 or img_region.shape[1] == 0:
        return False
    
    # 转换为灰度（如果不是）
    if len(img_region.shape) == 3:
        gray_region = cv2.cvtColor(img_region, cv2.COLOR_BGR2GRAY)
    else:
        gray_region = img_region
    
    # 检查1：高度不能太大（文字通常较矮）
    height_ratio = img_region.shape[0] / img_region.shape[1] if img_region.shape[1] > 0 else 1
    if height_ratio > 0.3:  # 如果高度超过宽度的30%，可能不是单行文字
        return False
    
    # 检查2：颜色是否偏暗（文字通常是黑色或深色）
    mean_brightness = np.mean(gray_region)
    if mean_brightness > 200:  # 太亮，不像文字
        return False
    
    # 检查3：内容像素占比（文字不会太密集也不会太稀疏）
    content_pixels = np.sum(thresh_region > 0)
    total_pixels = thresh_region.shape[0] * thresh_region.shape[1]
    density = content_pixels / total_pixels if total_pixels > 0 else 0
    
    if density < 0.05 or density > 0.5:  # 密度不在合理范围
        return False
    
    return True

def detect_and_remove_text(img, thresh, row_sums):
    """
    检测并移除图标上方或下方的文字标题
    
    返回: (top_crop, bottom_crop) - 需要裁剪的上下边界
    """
    h = len(row_sums)
    
    # 定义"空白行"的阈值（行和很小）
    empty_threshold = max(5, img.shape[1] * 0.01)  # 至少5，或宽度的1%
    # 定义"间隙"的最小行数
    min_gap_rows = max(2, int(h * 0.02))  # 至少2行，或高度的2%
    
    # 找到所有内容行（非空白行）
    content_rows = [i for i, val in enumerate(row_sums) if val > empty_threshold]
    
    if len(content_rows) == 0:
        return 0, h
    
    # 找到主要内容区域（最大的连续内容块）
    # 先找出所有的间隙
    gaps = []
    if len(content_rows) > 1:
        for i in range(len(content_rows) - 1):
            gap_size = content_rows[i + 1] - content_rows[i] - 1
            if gap_size >= min_gap_rows:
                gap_start = content_rows[i]
                gap_end = content_rows[i + 1]
                gaps.append((gap_start, gap_end, gap_size))
    
    top_crop = 0
    bottom_crop = h
    
    # 如果存在明显的间隙，说明可能有分离的文字
    if gaps:
        # 找到最大的间隙
        largest_gap = max(gaps, key=lambda x: x[2])
        gap_start, gap_end, gap_size = largest_gap
        
        # 计算间隙上方和下方的内容量和行数
        top_rows = gap_start
        bottom_rows = h - gap_end
        top_content = sum(row_sums[:gap_start])
        bottom_content = sum(row_sums[gap_end:])
        
        # 判断哪一部分是主要图标，哪一部分是文字
        # 文字的特征：1) 内容较少 2) 行数较少 3) 符合文字特征
        
        # 检查上方区域
        if top_rows > 0 and top_rows < h * 0.3:  # 上方行数不超过30%
            if top_content < bottom_content * 0.4:  # 上方内容明显少于下方
                # 进一步检查是否像文字
                top_region = img[:gap_start, :]
                top_thresh = thresh[:gap_start, :]
                if is_likely_text_region(top_region, top_thresh):
                    top_crop = gap_end
        
        # 检查下方区域
        if bottom_rows > 0 and bottom_rows < h * 0.3:  # 下方行数不超过30%
            if bottom_content < top_content * 0.4:  # 下方内容明显少于上方
                # 进一步检查是否像文字
                bottom_region = img[gap_end:, :]
                bottom_thresh = thresh[gap_end:, :]
                if is_likely_text_region(bottom_region, bottom_thresh):
                    bottom_crop = gap_start
    
    return top_crop, bottom_crop

def smart_crop_icon(img, padding=10):
    """
    单个 Icon 处理：去字、去空、加 Padding
    增强版：可以检测并删除上方或下方的文字标题
    """
    h, w = img.shape[:2]
    
    # 1. 裁剪掉可能残留的网格边缘 (比如四周切掉 3px)
    margin = 3
    if h > 2*margin and w > 2*margin:
        img = img[margin:-margin, margin:-margin]
        h, w = img.shape[:2]
    
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 2. 使用改进的方法检测并移除文字
    row_sums = np.sum(thresh, axis=1)
    top_crop, bottom_crop = detect_and_remove_text(img, thresh, row_sums)
    
    # 应用裁剪
    if top_crop > 0 or bottom_crop < h:
        img = img[top_crop:bottom_crop, :]
        thresh = thresh[top_crop:bottom_crop, :]

    # 3. 寻找 Icon 的精确边界
    coords = cv2.findNonZero(thresh)
    if coords is not None:
        x, y, w_box, h_box = cv2.boundingRect(coords)
        
        # 裁剪并添加 Padding
        # 创建一个新的白色画布
        final_h = h_box + 2 * padding
        final_w = w_box + 2 * padding
        canvas = np.ones((final_h, final_w, 3), dtype=np.uint8) * 255
        
        # 提取 icon 内容
        icon_content = img[y:y+h_box, x:x+w_box]
        
        # 将 icon 贴到画布中心
        canvas[padding:padding+h_box, padding:padding+w_box] = icon_content
        return canvas
    
    return img

def process_image_robust(image_path, labels_data):
    if not os.path.exists(image_path):
        print(f"Error: {image_path} not found.")
        return

    print(f"Processing: {image_path} ...")
    img = cv2.imread(image_path)
    
    # 1. 检测网格
    # 返回的是排序好的 (x, y, w, h) 列表
    grid_boxes = detect_grid_cells(img, expected_cols=6, expected_rows=4)
    
    # 2. 准备文本数据
    lines = [l.strip() for l in labels_data.strip().split('\n') if l.strip()]
    style = "flat"
    start_idx = 0
    if lines[0].lower().startswith("style:"):
        style = clean_filename(lines[0].split(':')[1])
        start_idx = 1
        
    output_dir = "extracted_icons"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 3. 遍历并保存
    for i, box in enumerate(grid_boxes):
        text_idx = start_idx + i
        if text_idx >= len(lines):
            break
            
        # 解析文本
        line_text = lines[text_idx]
        parts = line_text.split(',', 1)
        if len(parts) == 2:
            category = clean_filename(parts[0])
            name = clean_filename(parts[1])
        else:
            category = "icon"
            name = clean_filename(parts[0])
            
        filename = f"{category}-{name}-{style}.png"
        save_path = os.path.join(output_dir, filename)
        
        # 提取单元格
        x, y, w, h = box
        cell_img = img[y:y+h, x:x+w]
        
        # 智能裁切
        final_img = smart_crop_icon(cell_img, padding=10)
        
        cv2.imwrite(save_path, final_img)
        # print(f"Saved: {filename}") # 减少刷屏

    print(f"Done. Extracted {len(grid_boxes)} icons to '{output_dir}/'.\n")

def process_batch_range(start_batch, end_batch, base_dir="generated_icons"):
    """
    批量处理指定范围内的batch文件夹下的所有png文件
    
    Args:
        start_batch: 起始batch编号 (例如: 1)
        end_batch: 结束batch编号 (例如: 10)
        base_dir: batch文件夹所在的基础目录
    """
    print(f"开始批量处理 batch_{start_batch:04d} 到 batch_{end_batch:04d} ...\n")
    
    total_processed = 0
    failed_files = []
    
    for batch_num in range(start_batch, end_batch + 1):
        batch_dir = os.path.join(base_dir, f"batch_{batch_num:04d}")
        
        # 检查batch文件夹是否存在
        if not os.path.exists(batch_dir):
            print(f"Warning: {batch_dir} 不存在，跳过...")
            continue
        
        print(f"处理 {batch_dir} ...")
        
        # 查找该batch下的所有png文件
        png_files = glob.glob(os.path.join(batch_dir, "*.png"))
        
        if not png_files:
            print(f"  未找到png文件，跳过...")
            continue
        
        # 处理每个png文件
        for png_path in sorted(png_files):
            # 构造对应的txt文件路径
            txt_path = png_path.rsplit('.', 1)[0] + '.txt'
            
            # 检查txt文件是否存在
            if not os.path.exists(txt_path):
                print(f"  Warning: {txt_path} 不存在，跳过 {os.path.basename(png_path)}")
                failed_files.append(png_path)
                continue
            
            # 读取txt文件内容
            try:
                with open(txt_path, 'r', encoding='utf-8') as f:
                    labels_data = f.read()
                
                # 处理图像
                process_image_robust(png_path, labels_data)
                total_processed += 1
                
            except Exception as e:
                print(f"  Error processing {os.path.basename(png_path)}: {str(e)}")
                failed_files.append(png_path)
    
    # 输出总结
    print(f"\n{'='*60}")
    print(f"批量处理完成！")
    print(f"成功处理: {total_processed} 个文件")
    
    if failed_files:
        print(f"失败/跳过: {len(failed_files)} 个文件")
        print("失败文件列表:")
        for f in failed_files:
            print(f"  - {f}")
    print(f"{'='*60}")

if __name__ == "__main__":
    # 设置要处理的batch范围
    START_BATCH = 1    # 起始batch编号
    END_BATCH = 200     # 结束batch编号
    
    # 执行批量处理
    process_batch_range(START_BATCH, END_BATCH)