import cv2
import os
import numpy as np
import pytesseract

def locate_license_plate(image_path):
    if not os.path.isfile(image_path):
        print(f"Error: File '{image_path}' does not exist.")
        return None

    """
    根据输入的图像路径，定位蓝色车牌并返回裁剪出的车牌区域
    """
    # 读取图像
    car = cv2.imread(image_path, 1)

    if car is None:
        print(f"Error: Unable to read image '{image_path}'")
        return None

    # 定义蓝色所对应的色彩空间范围
    lower_blue = np.array([100, 110, 110])
    upper_blue = np.array([130, 255, 255])

    # 将图像转换到HSV颜色空间
    hsv = cv2.cvtColor(car, cv2.COLOR_BGR2HSV)

    # 获取蓝色区域的掩模
    mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)

    # 将掩模转换为灰度图像
    mask = mask_blue

    # 形态学处理，开运算和闭运算
    matrix = np.ones((20, 20), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, matrix)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, matrix)

    # 二值化
    ret, mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY)

    # 查找轮廓
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 寻找最大轮廓并定位车牌
    max_area = 0
    best_contour = None
    for contour in contours:
        area = cv2.contourArea(contour)
        if area > max_area:
            max_area = area
            best_contour = contour

    # 获取定位车牌的外接矩形框
    rect = cv2.minAreaRect(best_contour)
    box = cv2.boxPoints(rect)
    box = np.int32(box)

    # 获取旋转矩阵
    angle = rect[2]
    flag = 0
    print(angle)
    if angle < -45:
        angle += 90
        flag = 1
    if angle > 45:
        angle -= 90
        flag = 1
        print(angle)
    center = (rect[0][0], rect[0][1])
    size = (int(rect[1][0]), int(rect[1][1]))
    M = cv2.getRotationMatrix2D(center, angle, 1.0)

    # 旋转图像
    height, width = car.shape[:2]
    rotated = cv2.warpAffine(car, M, (width, height))

    # 获取旋转后矩形框的坐标
    if flag == 0:
        box = cv2.boxPoints(((center[0], center[1]), (size[0], size[1]), 0.0))
        box = np.int32(box)

    # 裁剪车牌区域
    xs = [box[0, 0], box[1, 0], box[2, 0], box[3, 0]]
    ys = [box[0, 1], box[1, 1], box[2, 1], box[3, 1]]
    x1, x2 = min(xs), max(xs)
    y1, y2 = min(ys), max(ys)
    ROI_plate = rotated[y1:y2, x1:x2]

    return ROI_plate

def preprocess_image_for_ocr(image):
    """
    预处理图像以提高 OCR 识别率
    """
    # # 去除左右上下各3个像素的边框
    height, width = image.shape[:2]
    image = image[3:height - 3, 3:width - 3]

    # 调整图像大小到 165x40
    image = cv2.resize(image, (165, 40), interpolation=cv2.INTER_AREA)

    # 转换为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 二值化
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # 去噪
    denoised = cv2.medianBlur(binary, 3)

    return denoised

def find_split_line(image):
    """
    在宽度15-30像素之间查找分割线
    """
    height, width = image.shape
    column_sums = np.sum(image[3:height - 3, 15:30], axis=0)  # 忽略顶部和底部的像素

    zero_columns = np.where(column_sums == 0)[0]  # 找到像素和为0的列
    if len(zero_columns) == 0:
        # 如果没有找到和为0的列，返回默认分割线为22
        split_line = 22
    else:
        # 计算零列的中值作为分割线
        split_line = int(np.median(zero_columns)) + 15  # 加上偏移量15
    print(split_line)
    return split_line

def recognize_characters(image):
    """
    使用OCR技术识别车牌上的字符
    """
    preprocessed_image = preprocess_image_for_ocr(image)

    # 找到分割线
    split_line = find_split_line(preprocessed_image)

    # 切割图像为两部分
    part1 = preprocessed_image[:, :split_line]  # 宽度0-split_line部分
    part2 = preprocessed_image[:, split_line:]  # 剩余部分

    # 调用 Tesseract OCR 识别第一部分
    custom_config_chi_sim = r'--oem 3 --psm 6 -l chi_sim'
    result_chi_sim = pytesseract.image_to_string(part1, config=custom_config_chi_sim)

    # 调用 Tesseract OCR 识别第二部分
    custom_config_default = r'--oem 3 --psm 6'
    result_default = pytesseract.image_to_string(part2, config=custom_config_default)

    # 合并两个部分的识别结果
    recognized_characters = result_chi_sim.strip() + result_default.strip()

    return recognized_characters