#!/usr/bin/env python3
"""
CCPD_Random_Blue5K_Green5K_EXP_CUT → ICE-COR 角点推理 → 中心外扩 → 透视变换
==================================================================================
输出: CCPD_Random_Blue5K_Green5K_EXP_CUT_TRANS/ (192×64, 文件名追加 _trans)

用法:
  python build_exp_cut_trans.py --full
  python build_exp_cut_trans.py --n 100
  python build_exp_cut_trans.py --dry-run
"""

import argparse
import os
import sys
import time
from pathlib import Path

import cv2
import numpy as np
import torch

SCRIPT_DIR  = Path(__file__).resolve().parent
PROJ_ROOT   = SCRIPT_DIR.parents[1]
ICE_COR_DIR = PROJ_ROOT / "ICE_NNv2" / "ICE-COR"
MODEL_PATH  = ICE_COR_DIR / "output" / "1_CCPD_19K_EXP_CUT" / "best.pth"
SRC_DIR     = PROJ_ROOT / "ICE_Dataset" / "CCPD_Random_Blue5K_Green5K_EXP_CUT"
DST_DIR     = PROJ_ROOT / "ICE_Dataset" / "CCPD_Random_Blue5K_Green5K_EXP_CUT_TRANS"

OUT_W, OUT_H = 192, 64
PAD_X = 0.08
PAD_Y = 0.05

if str(ICE_COR_DIR) not in sys.path:
    sys.path.insert(0, str(ICE_COR_DIR))
from corner_rectify import load_corner_model, predict_plate_corners, order_corners

# ---- Unicode-safe I/O ----
def _imread(path: str):
    with open(path, 'rb') as f:
        data = f.read()
    return cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)

def _imwrite(path: str, img: np.ndarray, params=None):
    ext = os.path.splitext(path)[1] or '.jpg'
    ok, buf = cv2.imencode(ext, img, params or [])
    if ok:
        with open(path, 'wb') as f:
            f.write(buf.tobytes())
    return ok

# ---- Single image ----
def process_one(model, device, img_path: Path, out_img_path: Path, draw: bool = False):
    img = _imread(str(img_path))
    if img is None:
        return "skip"
    H, W = img.shape[:2]
    try:
        corners = predict_plate_corners(img, model, device)
        if corners is None:
            return "skip"
    except Exception:
        return "skip"
    try:
        ordered = order_corners(corners)
    except Exception:
        return "skip"

    # center expand
    cx = ordered.mean(axis=0)
    expanded = ordered + (ordered - cx) * np.array([PAD_X, PAD_Y])
    expanded[:, 0] = np.clip(expanded[:, 0], 0, W - 1)
    expanded[:, 1] = np.clip(expanded[:, 1], 0, H - 1)

    src = expanded.astype(np.float32)
    dst = np.float32([[0, 0], [OUT_W-1, 0], [OUT_W-1, OUT_H-1], [0, OUT_H-1]])
    M = cv2.getPerspectiveTransform(src, dst)
    warped = cv2.warpPerspective(img, M, (OUT_W, OUT_H),
                                  flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)

    if draw:
        for k, (px, py) in enumerate(expanded):
            cv2.circle(warped, tuple(dst[k].astype(int)), 3, (0, 255, 0), -1)
        cv2.polylines(warped, [dst.astype(np.int32)], isClosed=True, color=(0, 255, 0), thickness=1)

    out_stem = img_path.stem + "_trans"
    out_img_path.with_stem(out_stem).parent.mkdir(parents=True, exist_ok=True)
    _imwrite(str(out_img_path.with_stem(out_stem)), warped, [cv2.IMWRITE_JPEG_QUALITY, 92])
    return "ok"

# ---- Main ----
def main():
    parser = argparse.ArgumentParser()
    g = parser.add_mutually_exclusive_group(required=True)
    g.add_argument('--full', action='store_true')
    g.add_argument('--n', type=int)
    g.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()

    print(f"Model: {MODEL_PATH}")
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model, device = load_corner_model(str(MODEL_PATH), device)
    model.eval()
    print(f"Device: {device}\n")

    classes = ['blue', 'green']
    draw = bool(args.dry_run)

    for cls in classes:
        src_cls = SRC_DIR / cls
        dst_cls = DST_DIR / cls
        dst_cls.mkdir(parents=True, exist_ok=True)
        images = sorted(src_cls.glob('*.jpg'))
        if not images:
            print(f"[WARN] {cls}: no images")
            continue

        n = 2 if args.dry_run else (args.n or len(images))
        images = images[:n]
        print(f"[{cls}] {len(images)} images...", flush=True)

        t0 = time.perf_counter()
        ok = skip = 0
        for img_path in images:
            out = dst_cls / img_path.name
            r = process_one(model, device, img_path, out, draw=draw)
            if r == "ok": ok += 1
            else: skip += 1
        elapsed = time.perf_counter() - t0
        total = ok + skip
        print(f"  → ok:{ok}/{total} skip:{skip}  ({elapsed:.1f}s, {total/elapsed:.1f} img/s)")
    print(f"\nDone. Output: {DST_DIR}")

if __name__ == '__main__':
    main()