"""
CCPD原始真图 去偏采样 + YOLO EXP_CUT + 角点TRANS处理
输出: ICE_Dataset/CCPD_REAL_TRANS/train/
省份均衡: 每省最多 MAX_PER_PROVINCE 张
管线: YOLO检出→12%/18%外扩→裁切→角点重映射→warp (与322K EXP_CUT一致)
"""
import os, sys, cv2, random, shutil
import numpy as np
from pathlib import Path
from collections import defaultdict

ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
OCR_TOOLS = ICEv2 / "ICE_NNv2/ICE-OCR/script/tools"
sys.path.insert(0, str(OCR_TOOLS))
from dataset import parse_ccpd_filename, imread_unicode

CCPD_DIRS = [
    ICEv2 / "ICE_Dataset/CCPD/CCPD2019",
    ICEv2 / "ICE_Dataset/CCPD/CCPD2020",
]
OUT_DIR = ICEv2 / "ICE_Dataset/CCPD_REAL_TRANS/train"
MAX_PER_PROVINCE = 300
SEED = 42

YOLO_PATH = str(ICEv2 / "ICE_NNv2/yolov5t/output/1_CCPD_19K/weights/best.pt")
EXPAND_NORMAL = 1.12    # 常规外扩 12%
EXPAND_SMALL  = 1.18    # 小车牌外扩 18%
SMALL_BOX_W = SMALL_BOX_H = 16
MIN_PAD = 2

random.seed(SEED)


def process_one(img_path, plate, model, img_size=640):
    """YOLO→外扩→裁切→角点重映射→warp (仿 EXP_CUT 管线)"""
    img = imread_unicode(str(img_path))
    if img is None:
        return None

    h, w = img.shape[:2]
    results = model(img_path, imgsz=img_size, verbose=False)
    if len(results) == 0 or results[0].boxes is None or len(results[0].boxes) == 0:
        return None

    boxes = results[0].boxes
    # 取置信度最高的框 (YOLO 输出 xyxy 格式)
    best_idx = boxes.conf.argmax()
    x1, y1, x2, y2 = boxes.xyxy[best_idx].cpu().numpy().astype(int)
    bw, bh = x2 - x1, y2 - y1

    # 外扩 (与 EXP_CUT 完全一致)
    expand_factor = EXPAND_SMALL if (bw < SMALL_BOX_W or bh < SMALL_BOX_H) else EXPAND_NORMAL
    pad_x = max(MIN_PAD, int(round(bw * (expand_factor - 1.0) / 2.0)))
    pad_y = max(MIN_PAD, int(round(bh * (expand_factor - 1.0) / 2.0)))

    cx1 = max(0, int(x1) - pad_x)
    cy1 = max(0, int(y1) - pad_y)
    cx2 = min(w, int(x2) + pad_x)
    cy2 = min(h, int(y2) + pad_y)
    if cx2 <= cx1 or cy2 <= cy1:
        return None

    cropped = img[cy1:cy2, cx1:cx2]

    # 角点: 用 YOLO 检出框的四角 (原始图坐标), 重映射到裁切图坐标
    corners_orig = np.float32([[x1, y1], [x2, y1], [x2, y2], [x1, y2]])  # TL, TR, BR, BL
    corners_crop = (corners_orig - np.array([cx1, cy1])).astype(np.float32)

    # 透视变换 → 192x64
    OW, OH = 192, 64
    dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
    M = cv2.getPerspectiveTransform(corners_crop, dst)
    warped = cv2.warpPerspective(cropped, M, (OW, OH))
    return warped


# ---- Step 0: 清空旧输出 ----
if OUT_DIR.exists():
    shutil.rmtree(str(OUT_DIR))
OUT_DIR.mkdir(parents=True, exist_ok=True)


# ---- Step 1: 收集CCPD真图, 按省份分组 ----
print("[1/3] Scanning CCPD directories...")
prov_files = defaultdict(list)  # province -> [(full_path, stem, plate)]
total = skipped = 0

for ccpd_dir in CCPD_DIRS:
    for root, dirs, files in os.walk(str(ccpd_dir)):
        for f in files:
            if not f.lower().endswith(".jpg"):
                continue
            indices, plate = parse_ccpd_filename(f)
            if indices is None:
                skipped += 1
                continue
            total += 1
            prov = plate[0]
            full_path = os.path.join(root, f)
            stem = Path(f).stem
            prov_files[prov].append((full_path, stem, plate))

print(f"  Valid: {total}, Skipped: {skipped}, Provinces: {len(prov_files)}")
for prov in sorted(prov_files, key=lambda p: -len(prov_files[p])):
    print(f"    {prov}: {len(prov_files[prov])}")


# ---- Step 2: 去偏采样 ----
print(f"\n[2/3] Sampling max {MAX_PER_PROVINCE} per province...")
sampled = []
for prov, items in prov_files.items():
    if len(items) <= MAX_PER_PROVINCE:
        sampled.extend(items)
    else:
        sampled.extend(random.sample(items, MAX_PER_PROVINCE))
random.shuffle(sampled)
print(f"  Sampled: {len(sampled)}")
prov_dist = defaultdict(int)
for _, _, plate in sampled:
    prov_dist[plate[0]] += 1
for prov in sorted(prov_dist):
    print(f"    {prov}: {prov_dist[prov]}")


# ---- Step 3: YOLO EXP_CUT + warp + 保存 ----
print(f"\n[3/3] YOLO EXP_CUT (12%/18%) + warp {len(sampled)} images -> {OUT_DIR}...")
print(f"  YOLO: {YOLO_PATH}")
from ultralytics import YOLO
model = YOLO(YOLO_PATH)
done = err = 0
for i, (full_path, stem, plate) in enumerate(sampled):
    try:
        warped = process_one(full_path, plate, model)
        if warped is None:
            err += 1; continue

        out_name = f"{stem}_real_blue_{plate}_trans.jpg"
        out_path = OUT_DIR / out_name

        ret, enc = cv2.imencode(".jpg", warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
        if ret:
            with open(str(out_path), "wb") as wf:
                wf.write(enc.tobytes())
            done += 1
        else:
            err += 1
    except Exception as e:
        err += 1
        if err <= 5:
            print(f"  [!] {Path(full_path).name}: {e}")

    if (i + 1) % 200 == 0:
        print(f"  [{i+1}/{len(sampled)}] done={done}, err={err}")

print(f"\nDone! Saved: {done}, Errors: {err}")
print(f"Output: {OUT_DIR} ({len(list(OUT_DIR.glob('*.jpg')))} files)")