#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
图像预处理工具
用于提高 OCR 识别准确率
"""

import cv2
import numpy as np


def preprocess_id_card_image(image_path, output_path=None):
    """
    预处理身份证图片以提高 OCR 识别率

    参数:
        image_path: 输入图片路径
        output_path: 输出图片路径（可选，如果提供则保存预处理后的图片）

    返回:
        预处理后的图片（numpy array）
    """
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")

    # 1. 转为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 2. 降噪
    denoised = cv2.fastNlMeansDenoising(gray, None, 10, 7, 21)

    # 3. 自适应二值化
    binary = cv2.adaptiveThreshold(
        denoised, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        11, 2
    )

    # 4. 形态学操作去除噪点
    kernel = np.ones((2, 2), np.uint8)
    morph = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

    # 5. 锐化（可选）
    kernel_sharpen = np.array([
        [-1, -1, -1],
        [-1,  9, -1],
        [-1, -1, -1]
    ])
    sharpened = cv2.filter2D(morph, -1, kernel_sharpen)

    # 保存预处理后的图片
    if output_path:
        cv2.imwrite(output_path, sharpened)

    return sharpened


def enhance_image_quality(image_path):
    """
    增强图片质量（对比度、亮度等）

    参数:
        image_path: 输入图片路径

    返回:
        增强后的图片（numpy array）
    """
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")

    # 转换到 LAB 色彩空间
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)

    # 对 L 通道进行直方图均衡化
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)

    # 合并通道
    enhanced_lab = cv2.merge((cl, a, b))
    enhanced = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2BGR)

    return enhanced


def auto_rotate_image(image_path):
    """
    自动旋转图片使文字方向正确

    参数:
        image_path: 输入图片路径

    返回:
        旋转后的图片（numpy array）
    """
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")

    # 转为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 霍夫变换检测直线
    lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)

    if lines is not None:
        # 计算平均角度
        angles = []
        for rho, theta in lines[:, 0]:
            angle = np.rad2deg(theta)
            angles.append(angle)

        # 获取主要角度
        median_angle = np.median(angles)

        # 旋转图片
        if abs(median_angle - 90) < 45:
            rotation_angle = median_angle - 90
        else:
            rotation_angle = median_angle

        # 执行旋转
        (h, w) = img.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, rotation_angle, 1.0)
        rotated = cv2.warpAffine(
            img, M, (w, h),
            flags=cv2.INTER_CUBIC,
            borderMode=cv2.BORDER_REPLICATE
        )

        return rotated

    return img


def resize_for_ocr(image_path, target_width=1500):
    """
    调整图片尺寸以适合 OCR

    参数:
        image_path: 输入图片路径
        target_width: 目标宽度（像素）

    返回:
        调整大小后的图片（numpy array）
    """
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")

    # 获取原始尺寸
    h, w = img.shape[:2]

    # 如果宽度小于目标宽度，则放大
    if w < target_width:
        ratio = target_width / w
        new_width = target_width
        new_height = int(h * ratio)
        resized = cv2.resize(
            img, (new_width, new_height),
            interpolation=cv2.INTER_CUBIC
        )
        return resized

    # 如果宽度大于目标宽度，则缩小
    elif w > target_width * 1.5:
        ratio = target_width / w
        new_width = target_width
        new_height = int(h * ratio)
        resized = cv2.resize(
            img, (new_width, new_height),
            interpolation=cv2.INTER_AREA
        )
        return resized

    return img


# 示例用法
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("使用方法: python image_preprocess.py <图片路径>")
        sys.exit(1)

    image_path = sys.argv[1]

    # 预处理
    processed = preprocess_id_card_image(image_path, "processed.jpg")
    print("预处理完成，保存为 processed.jpg")

    # 增强
    enhanced = enhance_image_quality(image_path)
    cv2.imwrite("enhanced.jpg", enhanced)
    print("质量增强完成，保存为 enhanced.jpg")

    # 调整大小
    resized = resize_for_ocr(image_path)
    cv2.imwrite("resized.jpg", resized)
    print("尺寸调整完成，保存为 resized.jpg")