"""
CBLPRD → CCPD_TRANS 数据集构建 (ICE-COR 矫正)
输出: ICE_Dataset/CBLPRD_TRANS/train/
"""
import sys, os, cv2, numpy as np, torch
from pathlib import Path
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Lock

ICEv2 = Path(r"C:\Users\Kirl\Desktop\ICEv2")
COR_DIR = ICEv2 / "ICE_NNv2/ICE-COR"
MODEL_COR = COR_DIR / "output/1_CCPD_19K_EXP_CUT/best.pth"
ZIP_PATH = r"C:\Users\Kirl\Downloads\CBLPRD-330k_v1.zip"
OUT_DIR = ICEv2 / "ICE_Dataset/CBLPRD_TRANS/train"

sys.path.insert(0, str(COR_DIR))
from corner_rectify import load_corner_model, predict_plate_corners, order_corners

# 类型→文件名标记
TYPE_MARKER = {
    '普通蓝牌': 'blue',
    '新能源小型车': 'green',
    '新能源大型车': 'new_energy',
    '单层黄牌': 'yellow',
    '黑色车牌': 'black',
}
SUPPORTED = set(TYPE_MARKER.keys())

print(f"COR: {MODEL_COR}")
dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cor_model, _ = load_corner_model(str(MODEL_COR), dev)

print(f"Output: {OUT_DIR}")
OUT_DIR.mkdir(parents=True, exist_ok=True)

# 读取 CBLPRD train.txt
zf = zipfile.ZipFile(ZIP_PATH, 'r')
labels = []
for line in zf.read('train.txt').decode('utf-8').strip().split('\n'):
    parts = line.strip().split(' ')
    if len(parts) < 3 or parts[2] not in SUPPORTED:
        continue
    fname = parts[0].split('/')[-1]
    labels.append((fname, parts[1], parts[2]))

# 同时加入 val.txt (合并训练)
for line in zf.read('val.txt').decode('utf-8').strip().split('\n'):
    parts = line.strip().split(' ')
    if len(parts) < 3 or parts[2] not in SUPPORTED:
        continue
    fname = parts[0].split('/')[-1]
    labels.append((fname, parts[1], parts[2]))

print(f"Total labels: {len(labels)}")
from collections import Counter
for t, c in Counter(l[2] for l in labels).most_common():
    print(f"  {t}: {c}")

# Build name→zip_path index
name_idx = {}
for n in zf.namelist():
    if n.endswith('.jpg'):
        name_idx[n.split('/')[-1]] = n
print(f"Zip index: {len(name_idx)} images")

# 统计已完成的
existing = set(f.stem.split('_trans')[0] if '_trans' in f.stem else '' for f in OUT_DIR.glob('*.jpg'))
print(f"Already done: {len(existing)}")

# 过滤已完成 (按原始文件名前缀)
pending = []
for fname, plate, ptype in labels:
    if fname in name_idx and fname.rsplit('.',1)[0] not in existing:
        pending.append((fname, plate, ptype, name_idx[fname]))

print(f"Pending: {len(pending)}")

lock = Lock()
count = [0]
total = len(pending)

def process_one(item):
    fname, plate, ptype, zip_path = item
    try:
        data = zf.read(zip_path)
        img = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            return f"READ_ERR:{fname}"

        # ICE-COR
        corners = predict_plate_corners(img, cor_model, dev)
        ordered = order_corners(corners)
        OW, OH = 192, 64
        cx, cy = ordered.mean(axis=0)
        exp = ordered + (ordered - np.array([[cx, cy]])) * np.array([0.08, 0.05])
        exp[:, 0] = np.clip(exp[:, 0], 0, img.shape[1] - 1)
        exp[:, 1] = np.clip(exp[:, 1], 0, img.shape[0] - 1)
        dst = np.float32([[0, 0], [OW - 1, 0], [OW - 1, OH - 1], [0, OH - 1]])
        M = cv2.getPerspectiveTransform(exp.astype(np.float32), dst)
        warped = cv2.warpPerspective(img, M, (OW, OH))
    except Exception as e:
        return f"COR_ERR:{fname}:{e}"

    # 文件名: {原始名}_{marker}_{plate}_trans.jpg
    marker = TYPE_MARKER[ptype]
    stem = fname.rsplit('.', 1)[0]
    out_name = f"{stem}_{marker}_{plate}_trans.jpg"
    out_path = OUT_DIR / out_name

    # Unicode safe write
    ret, enc = cv2.imencode('.jpg', warped, [cv2.IMWRITE_JPEG_QUALITY, 95])
    if ret:
        with open(str(out_path), 'wb') as f:
            f.write(enc.tobytes())
    else:
        return f"WRITE_ERR:{fname}"

    with lock:
        count[0] += 1
        if count[0] % 500 == 0:
            print(f"  [{count[0]}/{total}] {fname} -> {marker}:{plate}")
    return f"OK:{fname}"

# 多线程处理
WORKERS = 8
print(f"\nProcessing {total} images with {WORKERS} workers...")
with ThreadPoolExecutor(max_workers=WORKERS) as executor:
    futures = {executor.submit(process_one, item): item for item in pending}
    for future in as_completed(futures):
        result = future.result()
        if result.startswith("READ_ERR") or result.startswith("WRITE_ERR"):
            print(f"  [!] {result}")

zf.close()
print(f"\nDone! Total saved: {count[0]}/{total}")
print(f"Output: {OUT_DIR}")
print(f"Files: {len(list(OUT_DIR.glob('*.jpg')))}")