"""
CCPD原始真图 全量 TRANS (GT顶点外接框 + 12%/18%外扩)
输出: ICE_Dataset/CCPD_REAL_TRANS/train/
全量 ~338K, 文件名 _mask_real_ 标记 (训练时屏蔽省位+城市位)
"""
import os, sys, cv2, shutil
import numpy as np
from pathlib import Path

ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
sys.path.insert(0, str(ICEv2 / "ICE_NNv2/ICE-OCR/script/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"
EXPAND_NORMAL = 1.12
EXPAND_SMALL  = 1.18
SMALL_BOX_W = SMALL_BOX_H = 16
MIN_PAD = 2


def process_one(full_path, stem):
    """GT顶点→外接框→12/18%外扩→裁切→顶点重映射→warp"""
    img = imread_unicode(str(full_path))
    if img is None:
        return None
    h, w = img.shape[:2]

    parts = stem.split("-")
    if len(parts) < 4:
        return None
    # GT顶点 (CCPD: br, bl, tl, tr)
    verts = []
    for v in parts[3].split("_"):
        x, y = v.split("&")
        verts.append([float(x), float(y)])
    verts = np.float32(verts)

    # 外接框
    x1, y1 = verts[:, 0].min(), verts[:, 1].min()
    x2, y2 = verts[:, 0].max(), verts[:, 1].max()
    bw, bh = x2 - x1, y2 - y1

    # 外扩
    ef = EXPAND_SMALL if (bw < SMALL_BOX_W or bh < SMALL_BOX_H) else EXPAND_NORMAL
    px = max(MIN_PAD, int(round(bw * (ef - 1.0) / 2.0)))
    py = max(MIN_PAD, int(round(bh * (ef - 1.0) / 2.0)))
    cx1 = max(0, int(x1) - px); cy1 = max(0, int(y1) - py)
    cx2 = min(w, int(x2) + px); cy2 = min(h, int(y2) + py)
    if cx2 <= cx1 or cy2 <= cy1:
        return None

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

    # 顶点重映射到裁切坐标 + OpenCV顺序 (CCPD: br,bl,tl,tr → tl,tr,br,bl)
    verts_crop = verts[[2, 3, 0, 1], :] - np.array([cx1, cy1])
    verts_crop = verts_crop.astype(np.float32)

    OW, OH = 192, 64
    dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
    M = cv2.getPerspectiveTransform(verts_crop, dst)
    return cv2.warpPerspective(cropped, M, (OW, OH))


if __name__ == "__main__":
    # Step 0: 清空
    if OUT_DIR.exists():
        shutil.rmtree(str(OUT_DIR))
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    # Step 1: 收集
    print("[1/2] Scanning all CCPD...")
    all_files = []
    total = skipped = 0
    for d in CCPD_DIRS:
        for root, dirs, files in os.walk(str(d)):
            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
                all_files.append((os.path.join(root, f), Path(f).stem, plate))

    print(f"  Valid: {total}, Skipped: {skipped}")
    for prov in sorted(set(p[2][0] for p in all_files),
                       key=lambda p: -sum(1 for x in all_files if x[2][0] == p)):
        cnt = sum(1 for x in all_files if x[2][0] == prov)
        print(f"    {prov}: {cnt}")

    # Step 2: GT顶点 warp
    print(f"\n[2/2] GT-vertex warp {len(all_files)} -> {OUT_DIR}...")
    done = err = 0
    for i, (full_path, stem, plate) in enumerate(all_files):
        try:
            warped = process_one(full_path, stem)
            if warped is None:
                err += 1
                continue

            # _mask_real_ 标记: 训练时屏蔽 pos0(省)+pos1(城市码)
            out_name = f"{stem}_mask_real_{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 <= 3:
                print(f"  [!] {Path(full_path).name}: {e}")

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

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