import os

def generate_mmseg_datasets_my_dataset_file(output_file='./mmseg/datasets/my_dataset_model.py', dataset_class_name = "MyDataset_model", classes=['背景'], palette=[[0,0,0]], img_suffix=".png", seg_map_suffix="_gtFine_labelTrainIds.png", reduce_zero_label=False):
    # 先判断 classes 和 palette 的大小是否一致
    if len(classes) != len(palette):
        print(f"\033[91mInitial_Gen_mmseg_datasets_my_dataset.py程序中classes大小和palette大小不匹配，函数退出！\033[0m")
        return False
    
    # 判断是否有 '背景' 和 [0, 0, 0]
    if '背景' not in classes and 'background' not in classes and 'bg' not in classes and [0, 0, 0] not in palette:
        # 循环提示用户直到输入有效的值
        while True:
            print(f"现有Clas为：{classes}")
            print(f"现有palette为：{palette}")
            user_input = input("是否加入背景[0,0,0]到调色板第一位，Y加入，N不加入: ").strip().lower()
            if user_input == 'y':
                classes.insert(0,'背景')
                palette.insert(0,[0, 0, 0])
                print("已加入背景和调色板。")
                break
            elif user_input == 'n':
                print("未加入背景和调色板。")
                break
            else:
                print("无效输入，请输入Y或N。")
    
    # 使用 f-string 进行模板文本替换
    template = f'''# Copyright (c) OpenMMLab. All rights reserved.
# import mmengine.fileio as fileio

from mmseg.registry import DATASETS
from .basesegdataset import BaseSegDataset


@DATASETS.register_module()
class {dataset_class_name}(BaseSegDataset): # 表示你定义的数据的名字，顺便取一个名字即可
    """{dataset_class_name} dataset.
    """
    METAINFO = dict(
        classes={classes}, # 背景最好放到第一个
        palette={palette}) # TODO 标注类型和颜色

    def __init__(self,
                 img_suffix='{img_suffix}', # TODO mask图像类型
                 seg_map_suffix='{seg_map_suffix}', # TODO mask图像后缀
                 reduce_zero_label={reduce_zero_label}, # TODO 在第 0 类为无意义黑边时，使用reduce_zero_label = True将其和待分类内容分开；在第 0 类为 background 类别的数据集上，如果您最终是需要将背景和您的其余类别分开时，是不需要使用reduce_zero_label的 【reduce_zero_label = False】
                 **kwargs) -> None:
        super().__init__(
            img_suffix=img_suffix,
            seg_map_suffix=seg_map_suffix,
            reduce_zero_label=reduce_zero_label,
            **kwargs)
        
        # assert fileio.exists(
        #     self.data_prefix['img_path'], backend_args=self.backend_args)
'''

    # 将生成的程序内容写入文件
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(template)

    print(f"\033[93m{output_file} file generated successfully, 其中标签共{len(classes)}类\033[0m")
    return True, classes, palette

if __name__ == "__main__":

    ########### 定义各参数 ###########
    # 定义要替换的内容
    dataset_file_name = "my_dataset_model"
    dataset_class_name = "MyDataset_model"
    classes = ['肝脏','胆囊','分离钳','止血海绵','肝总管','胆总管','吸引器','剪刀','止血纱布','生物夹','无损伤钳','喷洒','胆囊管','胆囊动脉','电凝','标本袋','引流管','纱布','金属钛夹','术中超声','吻合器','乳胶管','推结器','肝带','钳夹','超声刀','脂肪','双极电凝','棉球','血管阻断夹','肿瘤','针','线','韧带','胆囊静脉']
    palette = [[255,91,0],[255,234,0],[85, 111, 181],[181, 227, 14],[72, 0, 255],[0, 155, 33],[255,0,255],[29, 32, 136],[160, 15, 95],[0,160,233],[52,184,178],[90,120,41],[255,0,0],[177,0,0],[167,24,233],[112,113,150],[0,255,0],[255,255,255],[0,255,255],[138,251,213],[136,162,196],[197,83,181],[202,202,200],[113,102,140],[66,115,82],[240,16,116],[155,132,0],[155,62,0],[146,175,236],[255,172,159],[245,161,0],[134,124,118],[0,157,142],[181,85,105],[42,8,66]]

    # 一般不太会变的参数
    img_suffix = ".png" 
    seg_map_suffix = "_gtFine_labelTrainIds.png"
    reduce_zero_label = False # 在第 0 类为 background 类别的数据集上，如果您最终是需要将背景和您的其余类别分开时，是不需要使用reduce_zero_label的 【reduce_zero_label = False】

    ########### 文件存储位置 ###########
    output_mmseg_datasets_dataset_file_name = os.path.join(f'./mmseg/datasets/{dataset_file_name}.py')

    # 生成程序文件
    success, classes, palette = generate_mmseg_datasets_my_dataset_file(output_file=output_mmseg_datasets_dataset_file_name, dataset_class_name=dataset_class_name, classes=classes, palette=palette, img_suffix=img_suffix, seg_map_suffix=seg_map_suffix, reduce_zero_label=reduce_zero_label)
