import os
import glob
import cv2
import numpy as np
import csv
import pytesseract
import re
from PIL import Image, ImageDraw
from pytesseract import Output

def crop_and_extract_first_word(image_path, height=45, min_blank=10):
    """
    切出图片顶部条带，精确分离第一个词
    :param image_path: 输入图片路径
    :param height: 顶部条带高度
    :param min_blank: 判断词结束的最小空白宽度
    :return: 第一个词的图像(PIL Image)
    """
    # 打开原始图片
    original_img = Image.open(image_path)
    width, img_height = original_img.size
    
    # 确保高度不超过图片高度
    crop_height = min(height, img_height)
    
    # 切出顶部条带
    top_strip = original_img.crop((0, 0, width, crop_height))
    
    # 使用OpenCV处理条带图像，分离第一个词
    strip_cv = cv2.cvtColor(np.array(top_strip), cv2.COLOR_RGB2BGR)
    
    # 增强图像预处理
    gray = cv2.cvtColor(strip_cv, cv2.COLOR_BGR2GRAY)
    
    # 1. 降噪处理
    denoised = cv2.fastNlMeansDenoising(gray, None, 10, 7, 21)
    
    # 2. 自适应阈值二值化（反色：文字为白色）
    binary = cv2.adaptiveThreshold(
        denoised, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV, 11, 7  # 优化参数
    )
    
    # 3. 形态学操作连接字符
    kernel = np.ones((2, 2), np.uint8)
    morphed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
    
    # 4. 按列统计非空白像素（文字部分）
    col_sum = np.sum(morphed, axis=0) / 255  # 每列文字像素数量
    
    # 5. 寻找第一个词的边界（改进算法）
    start_x = 0
    end_x = width - 1
    
    # 计算平均字符高度（用于过滤噪声）
    avg_char_height = np.mean(col_sum[col_sum > 0]) if np.any(col_sum > 0) else crop_height * 0.7
    
    # 寻找第一个有意义的字符起始位置
    for i in range(width - 5):
        # 连续5列有足够像素才认为是有效字符
        if all(col_sum[j] > avg_char_height * 0.3 for j in range(i, i+5)):
            start_x = i
            break
    
    # 寻找词结束位置
    if start_x < width - 5:
        blank_count = 0
        max_blank = min_blank
        
        # 从起始位置开始向右扫描
        for i in range(start_x, width):
            # 空白判断：列像素低于平均高度的20%
            if col_sum[i] < avg_char_height * 0.2:
                blank_count += 1
                if blank_count >= max_blank:
                    end_x = i - blank_count
                    break
            else:
                blank_count = 0
                
        # 如果没找到足够空白，使用备选方案
        if blank_count < max_blank:
            # 寻找第一个明显的下降点
            for i in range(start_x + 5, width - 5):
                if col_sum[i] < avg_char_height * 0.3 and col_sum[i+1] < avg_char_height * 0.3:
                    end_x = i
                    break
    
    # 安全边界检查
    start_x = max(0, start_x - 2)  # 向左扩展2像素
    end_x = min(width, end_x + 2)  # 向右扩展2像素
    
    # 确保至少提取一个字符
    if end_x <= start_x or (end_x - start_x) < 5:
        end_x = min(start_x + 50, width)  # 默认取50像素宽度
    
    # 提取第一个词并转换为PIL图像
    first_word_cv = strip_cv[:, start_x:end_x]
    
    # 保存中间结果用于调试
    cv2.imwrite("debug_binary.jpg", morphed)
    cv2.imwrite("debug_word.jpg", first_word_cv)
    
    first_word_pil = Image.fromarray(cv2.cvtColor(first_word_cv, cv2.COLOR_BGR2RGB))
    
    return first_word_pil

def ocr_word(image):
    """
    对单词图像进行OCR识别（法语）
    :param image: PIL图像对象
    :return: 处理后的OCR识别结果文本
    """
    # 转换为灰度图像并进行OCR（法语）
    gray = image.convert('L')
    
    # 增强OCR预处理
    gray_np = np.array(gray)
    gray_np = cv2.medianBlur(gray_np, 3)  # 中值滤波降噪
    _, thresh = cv2.threshold(gray_np, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    
    # 转换为PIL图像
    enhanced_img = Image.fromarray(thresh)
    
    # 使用法语OCR
    custom_config = r'--oem 3 --psm 8 -l fra'  # 法语OCR配置
    text = pytesseract.image_to_string(enhanced_img, config=custom_config).strip()
    
    # 处理OCR结果：去除开头数字和特殊字符
    processed_text = re.sub(r'^\d+', '', text)  # 移除开头的数字
    processed_text = re.sub(r'[^a-zA-ZÀ-ÿ\s]', '', processed_text)  # 只保留字母和法语特殊字符
    processed_text = processed_text.strip()  # 再次去除空白
    
    return processed_text

def process_folder_to_csv(input_folder, output_csv):
    """
    处理文件夹内所有JPG文件并保存结果到CSV
    :param input_folder: 输入文件夹路径
    :param output_csv: 输出CSV文件路径
    """
    # 获取所有JPG文件（按文件名排序）
    image_paths = sorted(glob.glob(os.path.join(input_folder, '*.jpg')))
    
    # 创建并写入CSV文件
    with open(output_csv, 'w', newline='', encoding='utf-8-sig') as csvfile:
        csv_writer = csv.writer(csvfile)
        csv_writer.writerow(['Filename', 'OCR Result'])  # 写入表头
        
        for img_path in image_paths:
            filename = os.path.basename(img_path)
            print(f"Processing: {filename}")
            
            try:
                # 提取第一个词
                word_image = crop_and_extract_first_word(img_path)
                
                # 进行OCR识别
                ocr_result = ocr_word(word_image)
                
                # 写入CSV
                csv_writer.writerow([filename, ocr_result])
                print(f"  OCR Result: {ocr_result}")
                
            except Exception as e:
                print(f"  Error processing {filename}: {str(e)}")
                csv_writer.writerow([filename, f"ERROR: {str(e)}"])

# 使用示例
if __name__ == "__main__":
    input_folder = r"D:\Dicts\Le_petit_Robert_1\output"  # 替换为你的图片文件夹路径
    output_csv = r"D:\Dicts\Le_petit_Robert_1\output\ocr_results.csv"      # 输出CSV文件名
    
    # 确保Tesseract路径正确（Windows系统需要设置）
    pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
    
    process_folder_to_csv(input_folder, output_csv)
    print(f"处理完成! 结果已保存到: {output_csv}")