import os
import cv2
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm

# ✅ 設定各資料夾的來源與輸出
datasets = [
    {
        "name": "train",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\stdc_724x362\train",
        "output_img_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\leftImg8bit\train",
        "output_mask_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\gtFine\train",
    },
    {
        "name": "val",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\stdc_724x362\valid",
        "output_img_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\leftImg8bit\val",
        "output_mask_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\gtFine\val",
    },
    {
        "name": "test",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\stdc_724x362\test",
        "output_img_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\leftImg8bit\test",
        "output_mask_dir": r"C:\Users\rd_de\kneronstdc\data\cityscapes\gtFine\test",
    }
]

# ✅ 建立 label ID 對應表（原始 -> 映射後）
label_mapping = {1: 0, 2: 1, 3: 2, 4: 3}

for dataset in datasets:
    name = dataset["name"]
    input_dir = dataset["input_dir"]
    output_img_dir = dataset["output_img_dir"]
    output_mask_dir = dataset["output_mask_dir"]

    os.makedirs(output_img_dir, exist_ok=True)
    os.makedirs(output_mask_dir, exist_ok=True)

    # ✅ 嘗試讀取 _classes.csv
    csv_path = os.path.join(input_dir, "_classes.csv")
    if not os.path.exists(csv_path):
        print(f"❌ 缺少 _classes.csv: {csv_path}")
        continue

    df = pd.read_csv(csv_path)
    valid_labels = set(df.iloc[:, 0].values)
    print(f"\n📂 資料集: {name}")
    print(f"✅ 合法類別 ID: {valid_labels}")

    for file in tqdm(os.listdir(input_dir), desc=f"轉換 {name}"):
        if not file.endswith("_mask.png"):
            continue

        mask_path = os.path.join(input_dir, file)
        image_name = file.replace("_mask.png", ".jpg")
        image_path = os.path.join(input_dir, image_name)

        if not os.path.exists(image_path):
            print(f"⚠️ 找不到對應圖片: {image_name}")
            continue

        img = cv2.imread(image_path)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        if img is None or mask is None:
            print(f"❌ 無法讀取圖像或 mask: {file}")
            continue

        labels = set(np.unique(mask))
        if not labels.issubset(valid_labels):
            print(f"❌ 非法類別出現在: {file}, labels={labels}")
            continue

        # ✅ 映射 label ID
        remapped_mask = np.zeros_like(mask)
        for old_id, new_id in label_mapping.items():
            remapped_mask[mask == old_id] = new_id

        # ✅ 命名符合 Cityscapes 格式
        stem = Path(file).stem.replace("_mask", "")
        out_img_path = os.path.join(output_img_dir, f"{stem}_leftImg8bit.png")
        out_mask_path = os.path.join(output_mask_dir, f"{stem}_gtFine_labelIds.png")

        cv2.imwrite(out_img_path, img)
        cv2.imwrite(out_mask_path, remapped_mask)

print("\n🎉 全部資料集轉換完成！Cityscapes 格式 OK")