import fitz  # PyMuPDF
import cv2
import numpy as np

def apply_vector_polka_dots_optimized(input_pdf_path, output_pdf_path, dot_radius=1, dot_spacing=5, text_padding=2, detect_dpi=400, debug_mode=True):
    try:
        doc = fitz.open(input_pdf_path)
    except Exception as e:
        print(f"无法打开 PDF: {e}")
        return

    print(f"开始处理 PDF (共 {len(doc)} 页)...")

    # --- 修改点 1: 放宽黄色检测范围 ---
    # 原来是 [20, 50, 50]，现在降低饱和度(S)和亮度(V)下限，以识别浅黄
    # 同时也稍微放宽色相(H)范围 [15-45]
    lower_yellow = np.array([15, 30, 30]) 
    upper_yellow = np.array([45, 255, 255])

    for page_num, page in enumerate(doc):
        zoom = detect_dpi / 72.0
        mat = fitz.Matrix(zoom, zoom)
        pix = page.get_pixmap(matrix=mat, alpha=False)
        
        img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, 3)
        img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        h_img, w_img = img_bgr.shape[:2]
        
        # --- 1. 生成黄色区域掩膜 ---
        hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
        mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
        
        # --- 修改点 2: 形态学膨胀 (填补细小空隙) ---
        # 如果某个点刚好落在黄色区域内的红线上，原来的代码会认为它不是黄色。
        # 膨胀操作会让黄色区域“吞噬”内部细小的杂色线条，保证填充连续。
        kernel = np.ones((3,3), np.uint8) # 3x3 的核
        mask_yellow = cv2.dilate(mask_yellow, kernel, iterations=2) # 膨胀2次

        if cv2.countNonZero(mask_yellow) == 0:
            print(f"  - 第 {page_num + 1} 页: 无黄色区域，跳过")
            continue

        # --- 2. 生成文字保护掩膜 ---
        mask_text_bboxes = np.zeros((h_img, w_img), np.uint8)
        
        # 建议：尝试使用 "words" 也许比 "blocks" 更精细，blocks 有时会包含很大空白
        # text_blocks = page.get_text("blocks") 
        text_words = page.get_text("words") # 改用 words 粒度更细
        
        padding_px = int(text_padding * zoom)

        # 使用 words 解析
        for word in text_words:
            x0_pdf, y0_pdf, x1_pdf, y1_pdf = word[:4] # words的前4个是坐标
            
            x0_px = int(x0_pdf * zoom)
            y0_px = int(y0_pdf * zoom)
            x1_px = int(x1_pdf * zoom)
            y1_px = int(y1_pdf * zoom)

            cv2.rectangle(mask_text_bboxes, 
                          (max(0, x0_px - padding_px), max(0, y0_px - padding_px)), 
                          (min(w_img, x1_px + padding_px), min(h_img, y1_px + padding_px)), 
                          255, -1)

        # --- 3. 计算最终绘制区域 ---
        mask_final = cv2.bitwise_and(mask_yellow, cv2.bitwise_not(mask_text_bboxes))

        # --- 修改点 3: 调试输出 (非常重要) ---
        if debug_mode:
            # 保存这一页的 mask 看看，白色区域就是要画点的地方
            # 如果你在图里看到原本是黄色的地方是黑的，说明HSV阈值还需要调
            debug_name = f"debug_mask_page_{page_num+1}.png"
            cv2.imwrite(debug_name, mask_final)
            print(f"    [Debug] 已保存掩膜图像: {debug_name} (请检查白色区域是否覆盖了所有黄色)")

        print(f"  - 第 {page_num + 1} 页: 正在写入矢量数据...")

        shape = page.new_shape()
        count = 0
        step_px = int(dot_spacing * zoom)
        if step_px < 1: step_px = 1

        for y in range(0, h_img, step_px):
            for x in range(0, w_img, step_px):
                if mask_final[y, x] > 0:
                    pdf_x = x / zoom
                    pdf_y = y / zoom
                    shape.draw_circle((pdf_x, pdf_y), dot_radius)
                    count += 1

        if count > 0:
            shape.finish(fill=(0, 0, 0), color=None)
            shape.commit(overlay=True) 

    doc.save(output_pdf_path, garbage=4, deflate=True)
    doc.close()
    print(f"\n处理完成！文件已保存: {output_pdf_path}")

if __name__ == "__main__":
    apply_vector_polka_dots_optimized(
        "input8.pdf", # 注意：这里你需要换回你的 input8.pdf
        "output_fixed.pdf", 
        dot_radius=1,
        dot_spacing=5,
        text_padding=1,
        detect_dpi=600,
        debug_mode=True # 开启调试模式
    )