import os
# 禁用所有可能导致 "could not execute a primitive" 错误的优化
os.environ["FLAGS_use_mkldnn"] = "0"
os.environ["FLAGS_prim_skip_onednn"] = "1"
os.environ["FLAGS_enable_pir_api"] = "0"
os.environ["FLAGS_use_cuda_managed_memory"] = "0"
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["OMP_NUM_THREADS"] = "1"

import cv2
import re
from paddleocr import PaddleOCR

# 初始化 PaddleOCR（只需初始化一次，可复用）
# 使用 CPU 模式并禁用可能的优化
ocr = PaddleOCR(
    use_angle_cls=True,
    lang='ch',
    use_gpu=False,
    enable_mkldnn=False,
    cpu_threads=1
)

def extract_text_from_image(image_path):
    """使用 PaddleOCR 从图片中提取文本"""
    try:
        # 使用 PaddleOCR 进行识别
        result = ocr.ocr(image_path)

        # 提取所有文本
        text_list = []
        if result and result[0]:
            for line in result[0]:
                text_list.append(line[1][0])

        # 合并所有识别的文本
        full_text = '\n'.join(text_list)
        return full_text.strip()
    except Exception as e:
        print(f"OCR 识别错误: {str(e)}")
        return ""

def extract_id_card_info(image_path):
    """
    从身份证图片中提取完整信息
    返回: {
        'name': 姓名,
        'gender': 性别,
        'nation': 民族,
        'birth': 出生日期,
        'address': 地址,
        'id_number': 身份证号
    }
    """
    try:
        # 使用 PaddleOCR 识别身份证
        result = ocr.ocr(image_path)

        name = None
        gender = None
        nation = None
        birth = None
        address = None
        id_number = None

        if result and result[0]:
            all_text = []
            for line in result[0]:
                text = line[1][0]
                confidence = line[1][1]
                all_text.append(text)

                # 识别身份证号（18位或15位）
                id_match = re.search(r'\d{17}[\dxX]|\d{15}', text)
                if id_match and confidence > 0.8:
                    id_number = id_match.group(0).upper()

                # 识别姓名
                if '姓名' in text:
                    name_match = re.search(r'姓名[:\s]*([^\s]+)', text)
                    if name_match:
                        name = name_match.group(1)
                elif not name and confidence > 0.9:
                    name_match = re.search(r'^[一-龥]{2,4}$', text)
                    if name_match and not any(kw in name_match.group(0) for kw in ['姓名', '性别', '民族', '出生', '住址', '公民']):
                        name = name_match.group(0)

                # 识别性别
                if '性别' in text:
                    gender_match = re.search(r'性别[:\s]*(男|女)', text)
                    if gender_match:
                        gender = gender_match.group(1)
                elif text in ['男', '女'] and confidence > 0.9:
                    gender = text

                # 识别民族
                if '民族' in text:
                    nation_match = re.search(r'民族[:\s]*([^\s]+)', text)
                    if nation_match:
                        nation = nation_match.group(1)
                elif confidence > 0.9 and re.match(r'^[一-龥]{1,4}族$', text):
                    nation = text

                # 识别出生日期
                if '出生' in text:
                    birth_match = re.search(r'(\d{4})年(\d{1,2})月(\d{1,2})日', text)
                    if birth_match:
                        birth = f"{birth_match.group(1)}-{birth_match.group(2).zfill(2)}-{birth_match.group(3).zfill(2)}"
                elif not birth:
                    birth_match = re.search(r'(\d{4})年(\d{1,2})月(\d{1,2})日', text)
                    if birth_match:
                        birth = f"{birth_match.group(1)}-{birth_match.group(2).zfill(2)}-{birth_match.group(3).zfill(2)}"

                # 识别住址
                if '住址' in text or '址' in text:
                    address_match = re.search(r'住?址[:\s]*(.+)', text)
                    if address_match:
                        address = address_match.group(1)

            # 如果住址没有完整识别，尝试合并多行文本
            if not address:
                full_text = '\n'.join(all_text)
                address_match = re.search(r'住?址[:\s]*([^\n]+(?:\n[^姓名性别民族出生公民身份号]+)*)', full_text)
                if address_match:
                    address = address_match.group(1).replace('\n', '')

        return {
            'name': name,
            'gender': gender,
            'nation': nation,
            'birth': birth,
            'address': address,
            'id_number': id_number
        }
    except Exception as e:
        print(f"身份证识别错误: {str(e)}")
        return {
            'name': None,
            'gender': None,
            'nation': None,
            'birth': None,
            'address': None,
            'id_number': None
        }

# 示例调用
if __name__ == "__main__":
    # 测试普通 OCR
    text = extract_text_from_image('path_to_image.png')
    print("提取的文本:", text)

    # 测试身份证识别
    id_info = extract_id_card_info('path_to_id_card.jpg')
    print("身份证信息:", id_info)