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

# ✅ 統一的訓練 label ID 對應表
name_to_unified_id = {
    'bunker': 0,
    'car': 1,
    'grass': 2,
    'greenery': 3,
    'person': 4,
    'pond': 5,
    'road': 6,
    'tree': 7,
    'background': 255  # 忽略背景
}

# ✅ 三組資料集，合併為 Cityscapes 格式的 train / val / test
datasets = [
    # 05-13
    {
        "name": "junior_train",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-13_danger_object_segmentation_junior\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": "junior_val",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-13_danger_object_segmentation_junior\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": "junior_test",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-13_danger_object_segmentation_junior\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",
    },

    # 05-19
    {
        "name": "may19_train",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-19_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": "may19_val",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-19_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": "may19_test",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-19_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",
    },

    # 05-09 ✅ 最新加入
    {
        "name": "may09_train",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-09_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": "may09_val",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-09_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": "may09_test",
        "input_dir": r"C:\Users\rd_de\kneronstdc\data\05-09_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",
    },
]

# ✅ 開始轉換
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)

    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)
    col_id, col_name = df.columns[0], df.columns[1]
    label_mapping = {
        int(row[col_id]): name_to_unified_id.get(str(row[col_name]).strip().lower(), 255)
        for _, row in df.iterrows()
    }

    print(f"\n📂 處理資料集: {name}")
    print(f"✅ Label 映射表: {label_mapping}")

    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)
        stem = file.replace("_mask.png", "")
        image_path = None
        for ext in [".jpg", ".png"]:
            candidate = os.path.join(input_dir, stem + ext)
            if os.path.exists(candidate):
                image_path = candidate
                break

        if image_path is None:
            print(f"⚠️ 找不到圖片: {stem}")
            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

        if img.shape[:2] != mask.shape:
            print(f"❌ 圖片與 mask 尺寸不一致: {file}")
            continue

        remapped_mask = np.full_like(mask, 255, dtype=np.uint8)
        for old_id, new_id in label_mapping.items():
            remapped_mask[mask == old_id] = new_id

        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(f"📊 {file} ➝ label 分布: {np.unique(remapped_mask)}")

print("\n✅ 所有資料集 train / val / test 合併完成，已轉換為 Cityscapes 格式！")