#!/usr/bin/env python3
"""
Build CCPD_19K dataset from:
  - CCPD2019 (all blue plates): random 5,000 images  -> class "blue"
  - CCPD2020 (all green plates): random 5,000 images -> class "green"
  - CCPD_SPE (9 special plate types): all 9,000 images -> 9 classes

Output structure (under ICE_Dataset/CCPD_19K/):
    train/
        blue/         ~4500
        green/        ~4500
        police/       ~900
        army/         ~900
        gangao/       ~900
        shiguan_old/  ~900
        shiguan_new/  ~900
        lingguan_old/ ~900
        lingguan_new/ ~900
        coach/        ~900
        new_energy/   ~900
    test/
        (same 11 classes, ~10% each)

Split ratio: 9:1 (train : test), stratified by class.
"""

import os
import sys
import random
import shutil
from pathlib import Path
from collections import defaultdict

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
BASE_DIR    = Path(r"C:\Users\Kirl\Desktop\ICEv2\ICE_Dataset")
CCPD2019_DIR = BASE_DIR / "CCPD" / "CCPD2019"
CCPD2020_DIR = BASE_DIR / "CCPD" / "CCPD2020"
CCPD_SPE_DIR = BASE_DIR / "CCPD_SPE"
OUTPUT_DIR  = BASE_DIR / "CCPD_19K"

SEED = 42
TRAIN_RATIO = 0.9  # 9:1 split

# Number of images to sample from CCPD2019 / CCPD2020
N_CCPD2019 = 5000
N_CCPD2020 = 5000

# ---------------------------------------------------------------------------
# Step 1: Collect all images from CCPD2019 (blue plates)
# ---------------------------------------------------------------------------
print("=" * 60)
print("Step 1: Collecting CCPD2019 (blue) images ...")

ccpd2019_images = []
for jpg in CCPD2019_DIR.rglob("*.jpg"):
    ccpd2019_images.append(jpg)

print(f"  Found {len(ccpd2019_images)} total images")
random.seed(SEED)
sampled_2019 = random.sample(ccpd2019_images, min(N_CCPD2019, len(ccpd2019_images)))
print(f"  Sampled {len(sampled_2019)} images")

# ---------------------------------------------------------------------------
# Step 2: Collect all images from CCPD2020 (green plates)
# ---------------------------------------------------------------------------
print("\nStep 2: Collecting CCPD2020 (green) images ...")

ccpd2020_images = []
for jpg in CCPD2020_DIR.rglob("*.jpg"):
    ccpd2020_images.append(jpg)

print(f"  Found {len(ccpd2020_images)} total images")
random.seed(SEED)
sampled_2020 = random.sample(ccpd2020_images, min(N_CCPD2020, len(ccpd2020_images)))
print(f"  Sampled {len(sampled_2020)} images")

# ---------------------------------------------------------------------------
# Step 3: Collect all images from CCPD_SPE (9 special classes)
# ---------------------------------------------------------------------------
print("\nStep 3: Collecting CCPD_SPE (9 special plate classes) ...")

# The 9 CCPD_SPE class names (subdirectory names)
SPE_CLASSES = [
    "police",
    "army",
    "gangao",
    "shiguan_old",
    "shiguan_new",
    "lingguan_old",
    "lingguan_new",
    "coach",
    "new_energy",
]

spe_images = defaultdict(list)
for jpg in CCPD_SPE_DIR.rglob("*.jpg"):
    # Path like: .../CCPD_SPE/ccpd_base/police/xxx.jpg
    # The class is the immediate parent directory name
    parent_dir = jpg.parent.name
    if parent_dir in SPE_CLASSES:
        spe_images[parent_dir].append(jpg)

for cls in SPE_CLASSES:
    print(f"  {cls}: {len(spe_images[cls])} images")

total_spe = sum(len(v) for v in spe_images.values())
print(f"  Total CCPD_SPE: {total_spe} images")

# ---------------------------------------------------------------------------
# Step 4: Build class -> file list mapping
# ---------------------------------------------------------------------------
print("\nStep 4: Building class mappings ...")

class_files = {}

# Blue plates from CCPD2019
class_files["blue"] = sampled_2019

# Green plates from CCPD2020
class_files["green"] = sampled_2020

# 9 special plate classes
for cls in SPE_CLASSES:
    class_files[cls] = spe_images[cls]

# Validate
all_11_classes = ["blue", "green"] + SPE_CLASSES
print(f"  11 classes: {all_11_classes}")
for cls in all_11_classes:
    print(f"  {cls}: {len(class_files[cls])} images")

# ---------------------------------------------------------------------------
# Step 5: Stratified 9:1 train/test split and copy
# ---------------------------------------------------------------------------
print("\nStep 5: Splitting and copying files ...")

random.seed(SEED)

for cls in all_11_classes:
    files = class_files[cls]
    random.shuffle(files)
    
    n_train = max(1, int(len(files) * TRAIN_RATIO))
    train_files = files[:n_train]
    test_files  = files[n_train:]
    
    # Create output directories
    train_dir = OUTPUT_DIR / "train" / cls
    test_dir  = OUTPUT_DIR / "test"  / cls
    train_dir.mkdir(parents=True, exist_ok=True)
    test_dir.mkdir(parents=True, exist_ok=True)
    
    # Copy train files (handle duplicate filenames by adding suffix)
    for src in train_files:
        dst = train_dir / src.name
        counter = 1
        stem, suffix = dst.stem, dst.suffix
        while dst.exists():
            dst = train_dir / f"{stem}_{counter}{suffix}"
            counter += 1
        shutil.copy2(src, dst)
    
    # Copy test files (handle duplicate filenames by adding suffix)
    for src in test_files:
        dst = test_dir / src.name
        counter = 1
        stem, suffix = dst.stem, dst.suffix
        while dst.exists():
            dst = test_dir / f"{stem}_{counter}{suffix}"
            counter += 1
        shutil.copy2(src, dst)
    
    print(f"  {cls}: train={len(train_files)}, test={len(test_files)}")

# ---------------------------------------------------------------------------
# Step 6: Summary
# ---------------------------------------------------------------------------
print("\n" + "=" * 60)
print("CCPD_19K dataset built successfully!")
print(f"  Output: {OUTPUT_DIR}")

total_train = 0
total_test  = 0
for cls in all_11_classes:
    train_cnt = len(list((OUTPUT_DIR / "train" / cls).glob("*.jpg")))
    test_cnt  = len(list((OUTPUT_DIR / "test"  / cls).glob("*.jpg")))
    total_train += train_cnt
    total_test  += test_cnt
    print(f"  {cls:15s}  train: {train_cnt:5d}  test: {test_cnt:5d}")

print(f"  {'TOTAL':15s}  train: {total_train:5d}  test: {total_test:5d}")
print(f"  Grand total: {total_train + total_test}")
print("=" * 60)