import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import argparse
import typing
import logging
from pathlib import Path
from dataclasses import dataclass

def extract_car_plate(img: cv.typing.MatLike) -> typing.Optional[cv.typing.MatLike]:
    """Extract the car plate part from given image.

    :param img: The image containing car plate in BGR format.
    :return: The image of binary car plate in U8 format if succeed, otherwise None.
    """
    # ── 1. 利用蓝色车牌颜色在 HSV 空间定位车牌 ──────────────────────────
    hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
    # 中国蓝牌 HSV 范围
    lower_blue = np.array([100, 80, 60])
    upper_blue = np.array([130, 255, 255])
    mask_blue = cv.inRange(hsv, lower_blue, upper_blue)

    # 形态学：闭运算填孔 + 开运算去噪
    kernel_close = cv.getStructuringElement(cv.MORPH_RECT, (25, 10))
    kernel_open  = cv.getStructuringElement(cv.MORPH_RECT, (5, 5))
    mask_blue = cv.morphologyEx(mask_blue, cv.MORPH_CLOSE, kernel_close)
    mask_blue = cv.morphologyEx(mask_blue, cv.MORPH_OPEN,  kernel_open)

    # ── 2. 连通域分析，筛选最符合车牌长宽比的区域 ──────────────────────
    num_labels, labels, stats, _ = cv.connectedComponentsWithStats(mask_blue, connectivity=8)

    best = None
    best_score = 0
    h_img, w_img = img.shape[:2]

    for i in range(1, num_labels):
        x, y, w, h, area = stats[i]
        if area < 3000:
            continue
        ratio = w / (h + 1e-5)
        # 标准车牌宽高比约 3:1 ~ 5:1
        if 2.5 < ratio < 6.0:
            score = area * (1 - abs(ratio - 3.5) / 3.5)
            if score > best_score:
                best_score = score
                best = (x, y, w, h)

    assert best is not None, "未找到车牌区域"
    x, y, w, h = best
    # 稍微扩边
    pad = 6
    x1 = max(x - pad, 0)
    y1 = max(y - pad, 0)
    x2 = min(x + w + pad, w_img)
    y2 = min(y + h + pad, h_img)

    plate_color = img[y1:y2, x1:x2].copy()
    print(f"车牌区域: x={x1}, y={y1}, w={x2-x1}, h={y2-y1}")

    # # 在原图上标记（仅供调试）
    # debug = img.copy()
    # cv.rectangle(debug, (x1, y1), (x2, y2), (0, 255, 0), 3)
    # cv.imwrite('./debug_detected.jpg', debug)

    # ── 3. 二值化：文字/边缘 → 黑色，背景 → 白色 ─────────────────────
    gray = cv.cvtColor(plate_color, cv.COLOR_BGR2GRAY)

    # 高斯模糊降噪
    blurred = cv.GaussianBlur(gray, (3, 3), 0)

    # Otsu 自动阈值，得到白字黑底，再取反 → 黑字白底
    _, binary_otsu = cv.threshold(blurred, 0, 255,
                                cv.THRESH_BINARY + cv.THRESH_OTSU)
    binary = cv.bitwise_not(binary_otsu)          # 反转：字符变黑，背景变白

    # 去除小噪点（开运算）
    kernel_denoise = cv.getStructuringElement(cv.MORPH_RECT, (2, 2))
    binary = cv.morphologyEx(binary, cv.MORPH_OPEN, kernel_denoise)

    return binary
    # cv.imwrite('./plate_binary.png', binary)
    # print("二值化结果已保存: plate_binary.png")

    # ── 4. 叠加边框轮廓（细化文字边缘，参考效果图）─────────────────────
    # Canny 边缘叠加让效果更接近参考图
    edges = cv.Canny(blurred, 40, 120)
    edges_inv = cv.bitwise_not(edges)             # 边缘→黑色
    combined = cv.bitwise_and(binary, edges_inv)  # 合并
    # 再做一次轻微腐蚀让字体略粗
    kernel_dilate = cv.getStructuringElement(cv.MORPH_RECT, (2, 2))
    combined = cv.erode(combined, kernel_dilate, iterations=1)

    return combined
    # cv.imwrite('./plate_final.png', combined)
    # print("最终结果已保存: plate_final.png")


@dataclass
class Cli:
    input_file: Path
    """The path to input file"""
    output_file: Path
    """The path to output file"""

    @staticmethod
    def from_cmdline() -> "Cli":
        # Build parser
        parser = argparse.ArgumentParser(
            prog="Car Plate Extractor",
            description="Extract the car plate part from given image.",
        )
        parser.add_argument(
            "-i",
            "--in",
            required=True,
            type=str,
            action="store",
            dest="input_file",
            metavar="in.jpg",
            help="""The path to input image containing car plate.""",
        )
        parser.add_argument(
            "-o",
            "--out",
            required=True,
            type=str,
            action="store",
            dest="output_file",
            metavar="out.png",
            help="""The path to output image for extracted car plate.""",
        )

        # Parse argument from cmdline and return
        args = parser.parse_args()
        return Cli(Path(args.input_file), Path(args.output_file))


def main():
    # Setup logging format
    logging.basicConfig(format="[%(levelname)s] %(message)s", level=logging.DEBUG)

    # Get user request
    cli = Cli.from_cmdline()

    # Load file
    in_img = cv.imread(str(cli.input_file), cv.IMREAD_COLOR)
    if in_img is None:
        logging.error(f"Fail to load image {cli.input_file}")
        return
    # Save extracted file if possible
    out_img = extract_car_plate(in_img)
    if out_img is not None:
        cv.imwrite(str(cli.output_file), out_img)


if __name__ == "__main__":
    main()