#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CCPD 真实车牌 ROI 提取工具
==========================
从 CCPD 图片中透视校正并裁剪车牌区域，保存为统一尺寸的 ROI 图片。
用于 CycleGAN/CUT 风格对齐训练的真实车牌域数据。

用法:
  python extract_real_plates.py --ccpd-dir ./CCPD/CCPD2019/ccpd_base --num 500 --output ./gan_data/real_plates

输出:
  440×140 的校正车牌 ROI (BGR .jpg)
"""

import os
import sys
import argparse
import random
from pathlib import Path
from typing import Optional, Tuple

import cv2
import numpy as np
from tqdm import tqdm


def parse_ccpd_corners(filename: str) -> Optional[np.ndarray]:
    """
    从 CCPD 文件名解析4角点，几何排序为 [TL, TR, BR, BL]。
    """
    base = Path(filename).stem
    parts = base.split('-')
    if len(parts) < 5:
        return None

    try:
        coords_str = parts[3]
        pts = []
        for pt in coords_str.split('_'):
            x, y = pt.split('&')
            pts.append((float(x), float(y)))

        if len(pts) != 4:
            return None

        a = np.array(pts)
        sums = a[:, 0] + a[:, 1]
        diffs = a[:, 0] - a[:, 1]

        tl = a[np.argmin(sums)]
        tr = a[np.argmax(diffs)]
        br = a[np.argmax(sums)]
        bl = a[np.argmin(diffs)]

        return np.float32([tl, tr, br, bl])
    except (ValueError, IndexError):
        return None


def extract_plate_roi(image_bgr: np.ndarray,
                       corners: np.ndarray,
                       output_size: Tuple[int, int] = (440, 140)) -> Optional[np.ndarray]:
    """
    透视校正提取车牌 ROI。
    """
    h, w = image_bgr.shape[:2]
    # 校验角点在图像范围内
    if corners.min() < -10 or corners[:, 0].max() > w + 10 or corners[:, 1].max() > h + 10:
        return None

    dst = np.float32([
        [0, 0],
        [output_size[0] - 1, 0],
        [output_size[0] - 1, output_size[1] - 1],
        [0, output_size[1] - 1],
    ])

    M = cv2.getPerspectiveTransform(corners, dst)
    roi = cv2.warpPerspective(image_bgr, M, output_size,
                               flags=cv2.INTER_CUBIC,
                               borderMode=cv2.BORDER_REPLICATE)
    return roi


def main():
    parser = argparse.ArgumentParser(description='CCPD 真实车牌 ROI 提取')
    parser.add_argument('--ccpd-dir', type=str, required=True,
                        help='CCPD 图片目录（多个用逗号分隔）')
    parser.add_argument('--num', type=int, default=500,
                        help='提取数量 (默认: 500)')
    parser.add_argument('--output', type=str, default='./gan_data/real_plates',
                        help='输出目录')
    parser.add_argument('--seed', type=int, default=42)

    args = parser.parse_args()
    random.seed(args.seed)
    np.random.seed(args.seed)

    # 收集图片
    ccpd_dirs = [d.strip() for d in args.ccpd_dir.split(',') if d.strip()]
    all_images = []
    for d in ccpd_dirs:
        dd = Path(d)
        if not dd.is_dir():
            print(f"[WARN] 目录不存在: {d}")
            continue
        for ext in ('*.jpg', '*.jpeg', '*.png'):
            all_images.extend(dd.glob(ext))

    random.shuffle(all_images)
    print(f"CCPD 图片总数: {len(all_images)}")
    print(f"目标提取数:   {args.num}")
    print(f"输出目录:     {args.output}")
    print("-" * 50)

    os.makedirs(args.output, exist_ok=True)

    success = 0
    for img_path in tqdm(all_images, desc='提取真实车牌ROI', unit='张'):
        if success >= args.num:
            break

        corners = parse_ccpd_corners(img_path.name)
        if corners is None:
            continue

        img = cv2.imread(str(img_path))
        if img is None:
            continue

        roi = extract_plate_roi(img, corners)
        if roi is None:
            continue

        out_name = f"real_{img_path.stem}.jpg"
        cv2.imwrite(os.path.join(args.output, out_name), roi,
                     [cv2.IMWRITE_JPEG_QUALITY, 95])
        success += 1

    print(f"\n完成! 提取: {success} 张")


if __name__ == '__main__':
    main()