# 322K 全自动流水线: 构建数据集 → EXP_CUT → TRANS → 训练
$ErrorActionPreference = "Continue"
$PY = "C:/Program Files/Python313/python.exe"
$ICE = "C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset"
$NN = "C:\Users\Kirl\Desktop\ICEv2\ICE_NNv2"

# ===== Step 1: Build CCPD_322K =====
Write-Host "`n===== 1/4 构建 CCPD_322K =====" -ForegroundColor Cyan

$SRC_SPE = "$ICE\CCPD_SPE\ccpd_base"
$SRC_RND = "$ICE\CCPD_Random_Blue150K_Green100K"
$DST = "$ICE\CCPD_322K"

if (Test-Path $DST) { Remove-Item -Recurse -Force $DST }
New-Item -ItemType Directory "$DST\train" -Force | Out-Null
New-Item -ItemType Directory "$DST\test" -Force | Out-Null

# SPE 9类 (9:1 split)
$types = @('police','army','gangao','shiguan_old','shiguan_new','lingguan_old','lingguan_new','coach','new_energy')
foreach ($t in $types) {
    $src = "$SRC_SPE\$t"
    if (-not (Test-Path $src)) { Write-Host "  SKIP $t"; continue }
    $files = Get-ChildItem $src -Filter *.jpg | Sort-Object Name
    $n = $files.Count
    $nTrain = [math]::Floor($n * 0.9)
    for ($i = 0; $i -lt $n; $i++) {
        $split = if ($i -lt $nTrain) { "train" } else { "test" }
        $d = "$DST\$split\$t"; New-Item -ItemType Directory $d -Force | Out-Null
        Copy-Item $files[$i].FullName "$d\" -Force
    }
    Write-Host "  $t : $n (train=$nTrain test=$($n-$nTrain))"
}

# Random蓝绿 (9:1 split)
foreach ($t in @('blue','green')) {
    $src = "$SRC_RND\$t"
    if (-not (Test-Path $src)) { Write-Host "  SKIP $t"; continue }
    $files = Get-ChildItem $src -Filter *.jpg | Sort-Object Name
    $n = $files.Count
    $nTrain = [math]::Floor($n * 0.9)
    for ($i = 0; $i -lt $n; $i++) {
        $split = if ($i -lt $nTrain) { "train" } else { "test" }
        $d = "$DST\$split\$t"; New-Item -ItemType Directory $d -Force | Out-Null
        Copy-Item $files[$i].FullName "$d\" -Force
    }
    Write-Host "  $t : $n (train=$nTrain test=$($n-$nTrain))"
}

$total = (Get-ChildItem $DST -Recurse -Filter *.jpg).Count
Write-Host "  TOTAL: $total" -ForegroundColor Green

# ===== Step 2: EXP_CUT =====
Write-Host "`n===== 2/4 YOLO EXP_CUT =====" -ForegroundColor Cyan
$EXP_CUT = "$ICE\CCPD_322K_EXP_CUT"
if (Test-Path $EXP_CUT) { Remove-Item -Recurse -Force $EXP_CUT }

# Patch the build script with our paths
$script = Get-Content "$ICE\CCPD_19K_EXP_CUT\build_ccpd_19k_exp_cut.py" -Raw -Encoding UTF8
$script = $script -replace 'SRC_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K"', "SRC_DIR = Path(r'$DST')"
$script = $script -replace 'DST_DIR = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT"', "DST_DIR = Path(r'$EXP_CUT')"
$tmp = "$ICE\scripts\_exp_cut_322k.py"
[System.IO.File]::WriteAllText($tmp, $script, [System.Text.UTF8Encoding]::new($false))
& $PY $tmp --full
Remove-Item $tmp

# ===== Step 3: TRANS =====
Write-Host "`n===== 3/4 ICE-COR TRANS =====" -ForegroundColor Cyan
$TRANS = "$ICE\CCPD_322K_EXP_CUT_TRANS"
if (Test-Path $TRANS) { Remove-Item -Recurse -Force $TRANS }

$script = Get-Content "$ICE\CCPD_19K_EXP_CUT_TRANS\build_ccpd_19k_exp_cut_trans.py" -Raw -Encoding UTF8
$script = $script -replace 'SRC_DIR     = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT"', "SRC_DIR     = Path(r'$EXP_CUT')"
$script = $script -replace 'DST_DIR     = PROJ_ROOT / "ICE_Dataset" / "CCPD_19K_EXP_CUT_TRANS"', "DST_DIR     = Path(r'$TRANS')"
$tmp = "$ICE\scripts\_trans_322k.py"
[System.IO.File]::WriteAllText($tmp, $script, [System.Text.UTF8Encoding]::new($false))
& $PY $tmp --full
Remove-Item $tmp

# ===== Step 4: Training =====
Write-Host "`n===== 4/4 ICE-OCR Training =====" -ForegroundColor Cyan
Write-Host "  Data: $TRANS"
Write-Host "  Launching..."
& $PY "$NN\ICE-OCR\script\train.py" --dataset ccpd_trans --epochs 100 --batch 128 --output ccpd_322k_best.pth --workers 2 --data-dir $TRANS

Write-Host "`n===== DONE =====" -ForegroundColor Green