#!/usr/bin/env python3
"""
C++ 管线批量重生 CCPD_SPE (9 类 x 7 子集 = 9,000 张)
比 Python 版快 ~5x，24 线程饱和
"""

import os
import sys
import time
import shutil
import subprocess
from pathlib import Path

ICE_DATASET = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
CPP_EXE = ICE_DATASET / 'CCPD_SPE_GEN' / 'cpp_pipeline' / 'build' / 'Release' / 'plate_swap.exe'
PLATE_POOL = ICE_DATASET / 'cpp_pipeline' / 'output' / 'plate_pool'
CCPD2019 = ICE_DATASET / 'CCPD' / 'CCPD2019'
SPE_OUT = ICE_DATASET / 'CCPD_SPE'

# 原始分布
DIST = {
    'ccpd_base':      588,
    'ccpd_blur':       55,
    'ccpd_challenge': 159,
    'ccpd_fn':         58,
    'ccpd_rotate':     30,
    'ccpd_tilt':       82,
    'ccpd_weather':    28,
}

TYPES = ['police','army','gangao','shiguan_old','shiguan_new',
         'lingguan_old','lingguan_new','coach','new_energy']

WORKERS = 24  # 物理核心数
SEED = 42

def main():
    if not CPP_EXE.exists():
        print(f"[FATAL] C++ exe not found: {CPP_EXE}")
        return 1
    if not PLATE_POOL.is_dir():
        print(f"[FATAL] Plate pool not found: {PLATE_POOL}")
        return 1

    # 清空旧数据
    if SPE_OUT.exists():
        print(f"删除旧 CCPD_SPE: {SPE_OUT}")
        shutil.rmtree(str(SPE_OUT))

    total_start = time.time()
    total = 0

    for plate_type in TYPES:
        t0 = time.time()
        type_count = 0
        print(f"\n{'='*50}")
        print(f"[{plate_type}]")
        print(f"{'='*50}")

        for subdir, count in DIST.items():
            ccpd_dir = CCPD2019 / subdir
            if not ccpd_dir.is_dir():
                continue

            out_dir = SPE_OUT / subdir / plate_type
            out_dir.mkdir(parents=True, exist_ok=True)

            cmd = [
                str(CPP_EXE),
                '--ccpd-dir', str(ccpd_dir),
                '--plate-type', plate_type,
                '--plate-pool', str(PLATE_POOL),
                '--output', str(out_dir),
                '--num', str(count),
                '--workers', str(WORKERS),
                '--seed', str(SEED + hash(plate_type) % 10000),
            ]

            print(f"  {subdir}: {count} 张...", end=' ', flush=True)
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode == 0:
                print(f"OK")
                type_count += count
            else:
                print(f"FAIL")
                print(f"    stderr: {result.stderr[:200]}")

        elapsed = time.time() - t0
        print(f"  [{plate_type}] {type_count} 张, {elapsed:.0f}s")
        total += type_count

    total_t = time.time() - total_start
    print(f"\n{'='*50}")
    print(f"DONE! {total} 张, {total_t/60:.1f}min")
    print(f"{'='*50}")

if __name__ == '__main__':
    os.chdir(str(ICE_DATASET))
    main()