# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import mmcv
import numpy as np
from mmcv.utils import print_log
from PIL import Image

from .builder import DATASETS
from .custom import CustomDataset

@DATASETS.register_module()
class GolfDataset(CustomDataset):
    """GolfDataset for semantic segmentation with four classes: car, grass, people, and road."""

    # ✅ 固定的類別與調色盤（不從 config 接收）
    CLASSES = ('car', 'grass', 'people', 'road')
    PALETTE = [
        [246, 14, 135],   # car
        [233, 81, 78],    # grass
        [220, 148, 21],   # people
        [207, 215, 220],  # road
    ]

    def __init__(self,
                 img_suffix='_leftImg8bit.png',
                 seg_map_suffix='_gtFine_labelIds.png',
                 **kwargs):
        super(GolfDataset, self).__init__(
            img_suffix=img_suffix,
            seg_map_suffix=seg_map_suffix,
            **kwargs)

        # ✅ DEBUG：初始化時印出 CLASSES 與 PALETTE
        print("✅ [GolfDataset] 初始化完成")
        print(f"   ➤ CLASSES: {self.CLASSES}")
        print(f"   ➤ PALETTE: {self.PALETTE}")
        print(f"   ➤ img_suffix: {img_suffix}")
        print(f"   ➤ seg_map_suffix: {seg_map_suffix}")
        print(f"   ➤ img_dir: {self.img_dir}")
        print(f"   ➤ ann_dir: {self.ann_dir}")
        print(f"   ➤ dataset length: {len(self)}")

    def results2img(self, results, imgfile_prefix, indices=None):
        """Write the segmentation results to images."""
        if indices is None:
            indices = list(range(len(self)))

        mmcv.mkdir_or_exist(imgfile_prefix)
        result_files = []
        for result, idx in zip(results, indices):
            filename = self.img_infos[idx]['filename']
            basename = osp.splitext(osp.basename(filename))[0]
            png_filename = osp.join(imgfile_prefix, f'{basename}.png')

            result = result.astype(np.uint8)

            # ✅ 把所有無效類別設為 255（當作背景處理）
            result[result >= len(self.PALETTE)] = 255

            output = Image.fromarray(result).convert('P')

            # ✅ 建立 palette，支援背景 class 255 為黑色
            palette = np.zeros((256, 3), dtype=np.uint8)
            for label_id, color in enumerate(self.PALETTE):
                palette[label_id] = color
            palette[255] = [0, 0, 0]  # 黑色背景

            output.putpalette(palette)
            output.save(png_filename)
            result_files.append(png_filename)

        return result_files

    def format_results(self, results, imgfile_prefix, indices=None):
        """Format the results into dir (for evaluation or visualization)."""
        return self.results2img(results, imgfile_prefix, indices)

    def evaluate(self,
                 results,
                 metric='mIoU',
                 logger=None,
                 imgfile_prefix=None):
        """Evaluate the results with the given metric."""

        # ✅ DEBUG：評估時印出目前 CLASSES 使用狀況
        print("🧪 [GolfDataset.evaluate] 被呼叫")
        print(f"   ➤ 當前 CLASSES: {self.CLASSES}")
        print(f"   ➤ 評估 metric: {metric}")
        print(f"   ➤ 結果數量: {len(results)}")

        metrics = metric if isinstance(metric, list) else [metric]
        eval_results = super(GolfDataset, self).evaluate(results, metrics, logger)

        # ✅ DEBUG：印出最終的 eval_results keys
        print(f"   ➤ 返回評估指標: {list(eval_results.keys())}")
        return eval_results