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\06-03_danger_object_segmentation\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\06-03_danger_object_segmentation\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\06-03_danger_object_segmentation\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）
label_mapping = {
    1: 0,  # bunker
    2: 1,  # car
    3: 2,  # grass
    4: 3,  # greenery
    5: 4,  # person
    6: 5,  # road
    7: 6   # tree
}

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))
        valid_ids = set(label_mapping.keys()) | {0}
        if not labels.issubset(valid_ids):
            print(f"❌ 非法類別出現在: {file}, labels={labels}")
            continue

        # ✅ 映射 label，0 ➝ 255（ignore），其餘照 mapping
        remapped_mask = np.full_like(mask, 255, dtype=np.uint8)  # 預設全部是 ignore
        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")