#!/usr/bin/env python3
"""
全量生成 Random 蓝绿车牌
blue 150K + green 100K, 覆盖全部 CCPD 背景
"""

import os, sys, time
from pathlib import Path

ICE_DATASET = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
os.chdir(str(ICE_DATASET))

sys.path.insert(0, str(ICE_DATASET / 'CCPD_SPE_GEN' / 'cv_pipeline'))
from plate_swap import PlateSwapPipeline

CCPD2019 = ICE_DATASET / 'CCPD' / 'CCPD2019'
CCPD2020 = ICE_DATASET / 'CCPD' / 'CCPD2020'
OUT = ICE_DATASET / 'CCPD_Random_Blue150K_Green100K'

# 全部可用 CCPD 子集 (含全量, 不限于2019)
CCPD_DIRS = [
    str(CCPD2019 / 'ccpd_base'),
    str(CCPD2019 / 'ccpd_blur'),
    str(CCPD2019 / 'ccpd_challenge'),
    str(CCPD2019 / 'ccpd_fn'),
    str(CCPD2019 / 'ccpd_rotate'),
    str(CCPD2019 / 'ccpd_tilt'),
    str(CCPD2019 / 'ccpd_weather'),
]

# 只保留实际存在的
CCPD_DIRS = [d for d in CCPD_DIRS if Path(d).is_dir()]
print(f"CCPD 背景目录: {len(CCPD_DIRS)}")

SEED = 42
WORKERS = 24

def main():
    import shutil
    # 断点续传: 不删除已有数据
    total_start = time.time()

    for plate_type, target in [('blue', 150000), ('green', 100000)]:
        out_dir = OUT / plate_type
        out_dir.mkdir(parents=True, exist_ok=True)
        existing = len(list(out_dir.glob('*.jpg')))
        if existing >= target:
            print(f"\n[{plate_type}] 已有 {existing} >= {target}, 跳过")
            continue
        need = target - existing
        print(f"\n{'='*60}")
        print(f"[{plate_type}] 已有 {existing}, 需补 {need} → 目标 {target}")
        print(f"{'='*60}")

        pipeline = PlateSwapPipeline(plate_type=plate_type)

        print(f"  生成中...", end=' ', flush=True)
        t0 = time.time()
        pipeline.process_batch(
            ccpd_dirs=CCPD_DIRS,
            output_dir=str(out_dir),
            num_images=need,
            preview=False,
            seed=SEED,
            workers=WORKERS,
        )
        elapsed = time.time() - t0
        print(f"完成 ({elapsed:.0f}s, {need/elapsed:.0f} 张/s)")

    total_t = time.time() - total_start
    blue_final = len(list((OUT / 'blue').glob('*.jpg')))
    green_final = len(list((OUT / 'green').glob('*.jpg')))
    print(f"\nDONE! blue={blue_final} green={green_final}, {total_t/60:.1f}min")

if __name__ == '__main__':
    main()