#!/usr/bin/env python3
"""322K Pipeline Step 2-4: EXP_CUT -> TRANS -> Train"""
import os, sys, shutil, subprocess
from pathlib import Path

ICE = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset')
PY = r'C:/Program Files/Python313/python.exe'
NN = Path(r'C:\Users\Kirl\Desktop\ICEv2\ICE_NNv2')

DATASET = ICE / 'CCPD_322K'
EXP_CUT = ICE / 'CCPD_322K_EXP_CUT'
TRANS   = ICE / 'CCPD_322K_EXP_CUT_TRANS'

def run(cmd):
    print(f"  {str(cmd)[:150]}...")
    if isinstance(cmd, list):
        r = subprocess.run(cmd)
    else:
        r = subprocess.run(cmd, shell=True)
    if r.returncode != 0: print(f"  FAILED ({r.returncode})"); sys.exit(1)

# ===== Step 2: EXP_CUT =====
print("===== 2/4 EXP_CUT =====")
shutil.rmtree(EXP_CUT, ignore_errors=True)
EXP_CUT.mkdir(parents=True)

exp_py = ICE / 'CCPD_19K_EXP_CUT' / 'build_ccpd_19k_exp_cut.py'
code = exp_py.read_text(encoding='utf-8')
code = code.replace(
    'SRC_DIR     = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K"          # 源数据集',
    f'SRC_DIR     = Path(r"{DATASET}")')
code = code.replace(
    'DST_DIR     = SCRIPT_DIR                                        # 本目录 (CCPD_19K_EXP_CUT)',
    f'DST_DIR     = Path(r"{EXP_CUT}")')
tmp = ICE / 'scripts' / '_e.py'
tmp.write_text(code, encoding='utf-8')
run([PY, str(tmp), '--full'])
tmp.unlink()
print(f"  Images: {sum(1 for _ in EXP_CUT.rglob('*.jpg'))}")

# ===== Step 3: TRANS =====
print("===== 3/4 TRANS =====")
shutil.rmtree(TRANS, ignore_errors=True)
TRANS.mkdir(parents=True)

trans_py = ICE / 'CCPD_19K_EXP_CUT_TRANS' / 'build_ccpd_19k_exp_cut_trans.py'
code = trans_py.read_text(encoding='utf-8')
code = code.replace(
    'SRC_DIR     = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT"         # 源数据集',
    f'SRC_DIR     = Path(r"{EXP_CUT}")')
code = code.replace(
    'DST_DIR     = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT_TRANS"',
    f'DST_DIR     = Path(r"{TRANS}")')
tmp = ICE / 'scripts' / '_t.py'
tmp.write_text(code, encoding='utf-8')
run([PY, str(tmp), '--full'])
tmp.unlink()
print(f"  Images: {sum(1 for _ in TRANS.rglob('*.jpg'))}")

# ===== Step 4: Training =====
print("===== 4/4 Training =====")
os.chdir(str(NN / 'ICE-OCR'))
run([PY, 'script/train.py', '--dataset', 'ccpd_trans', '--epochs', '100', '--batch', '128', '--output', 'ccpd_322k_best.pth', '--workers', '2', '--data-dir', str(TRANS)])
print("DONE")