#!/usr/bin/env python3
"""
验证码识别脚本
输入：验证码图片路径
输出：识别结果（打印到 stdout）

预处理流程：
1. 灰度化
2. 去除水平干扰线（形态学开运算）
3. 二值化
4. 去噪
5. ddddocr 识别
"""
import sys
import cv2
import numpy as np
import ddddocr

def preprocess(img_path):
    """多种预处理方案"""
    img = cv2.imread(img_path)
    if img is None:
        return []

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    results = []

    # 方案1：去水平线 + 二值化
    # 用水平核做形态学开运算，提取水平线
    h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 1))
    h_lines = cv2.morphologyEx(gray, cv2.MORPH_OPEN, h_kernel)
    # 从原图减去水平线
    no_lines = cv2.subtract(gray, h_lines)
    # 二值化
    _, binary1 = cv2.threshold(no_lines, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    results.append(binary1)

    # 方案2：自适应阈值
    adaptive = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                      cv2.THRESH_BINARY, 11, 2)
    results.append(adaptive)

    # 方案3：去水平线 + 自适应阈值
    adaptive2 = cv2.adaptiveThreshold(no_lines, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                       cv2.THRESH_BINARY, 11, 2)
    results.append(adaptive2)

    # 方案4：反色 + 去线 + 二值化
    inverted = cv2.bitwise_not(gray)
    h_lines_inv = cv2.morphologyEx(inverted, cv2.MORPH_OPEN, h_kernel)
    no_lines_inv = cv2.subtract(inverted, h_lines_inv)
    _, binary4 = cv2.threshold(no_lines_inv, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    results.append(binary4)

    # 方案5：中值滤波去噪 + OTSU
    median = cv2.medianBlur(gray, 3)
    _, binary5 = cv2.threshold(median, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    results.append(binary5)

    # 方案6：放大 + 去线 + 锐化
    big = cv2.resize(gray, None, fx=3, fy=3, interpolation=cv2.INTER_CUBIC)
    h_kernel_big = cv2.getStructuringElement(cv2.MORPH_RECT, (75, 1))
    h_lines_big = cv2.morphologyEx(big, cv2.MORPH_OPEN, h_kernel_big)
    no_lines_big = cv2.subtract(big, h_lines_big)
    _, binary6 = cv2.threshold(no_lines_big, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # 去小噪点
    kernel_small = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
    binary6 = cv2.morphologyEx(binary6, cv2.MORPH_OPEN, kernel_small)
    results.append(binary6)

    # 方案7：原图直接用
    results.append(gray)

    # 方案8: 原始彩色图（ddddocr 有时对彩色图效果更好）
    results.append(img)

    return results


def main():
    if len(sys.argv) < 2:
        print("用法: python3 ocr_captcha.py <图片路径>", file=sys.stderr)
        sys.exit(1)

    img_path = sys.argv[1]

    ocr = ddddocr.DdddOcr(show_ad=False)

    # 先直接用原图试
    with open(img_path, 'rb') as f:
        raw_result = ocr.classification(f.read())

    if raw_result and len(raw_result) == 4 and raw_result.isalnum():
        print(raw_result)
        return

    # 多种预处理方案
    preprocessed = preprocess(img_path)

    best = raw_result
    for i, processed in enumerate(preprocessed):
        try:
            if len(processed.shape) == 3:
                _, buf = cv2.imencode('.png', processed)
            else:
                _, buf = cv2.imencode('.png', processed)
            result = ocr.classification(buf.tobytes())

            if result and len(result) == 4 and result.isalnum():
                print(result)
                return

            # 记录最佳结果
            if result and len(result) >= 3 and (not best or len(result) == 4):
                best = result
        except Exception:
            continue

    # 返回最佳结果
    if best:
        print(best[:4] if len(best) > 4 else best)
    else:
        print("FAIL", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()