import os

def generate_mmseg_datasets_init_file(output_file, dataset_file_names=["my_data_set_model"], dataset_class_names=["MyDataset_model"]):
    # 判断数据集文件名 和 对应类名大小是否相同 # 不同则输出错误内容，并退出
    if len(dataset_file_names) != len(dataset_class_names):
        print(f"\033[91mInitial_Gen_mmseg_datasets_init_.py程序中MyDataset_File_name大小和MyDataset_Class_name大小不匹配，函数退出！\033[0m")
        return False

    # 模板开头部分
    header = """# Copyright (c) OpenMMLab. All rights reserved.
# yapf: disable
from .ade import ADE20KDataset
from .basesegdataset import BaseCDDataset, BaseSegDataset
from .bdd100k import BDD100KDataset
from .chase_db1 import ChaseDB1Dataset
from .cityscapes import CityscapesDataset
from .coco_stuff import COCOStuffDataset
from .dark_zurich import DarkZurichDataset
from .dataset_wrappers import MultiImageMixDataset
from .decathlon import DecathlonDataset
from .drive import DRIVEDataset
from .dsdl import DSDLSegDataset
from .hrf import HRFDataset
from .hsi_drive import HSIDrive20Dataset
from .isaid import iSAIDDataset
from .isprs import ISPRSDataset
from .levir import LEVIRCDDataset
from .lip import LIPDataset
from .loveda import LoveDADataset
from .mapillary import MapillaryDataset_v1, MapillaryDataset_v2
from .night_driving import NightDrivingDataset
from .nyu import NYUDataset
from .pascal_context import PascalContextDataset, PascalContextDataset59
from .potsdam import PotsdamDataset
from .refuge import REFUGEDataset
from .stare import STAREDataset
from .synapse import SynapseDataset
"""
    # 增加多个 dataset_file_names imports
    imports = ""
    for dataset_file, dataset_name in zip(dataset_file_names, dataset_class_names):
        imports += f"from .{dataset_file} import {dataset_name} # TODO\n"
    
    # 中间固定部分
    middle = """# yapf: disable
from .transforms import (CLAHE, AdjustGamma, Albu, BioMedical3DPad,
                         BioMedical3DRandomCrop, BioMedical3DRandomFlip,
                         BioMedicalGaussianBlur, BioMedicalGaussianNoise,
                         BioMedicalRandomGamma, ConcatCDInput, GenerateEdge,
                         LoadAnnotations, LoadBiomedicalAnnotation,
                         LoadBiomedicalData, LoadBiomedicalImageFromFile,
                         LoadImageFromNDArray, LoadMultipleRSImageFromFile,
                         LoadSingleRSImageFromFile, PackSegInputs,
                         PhotoMetricDistortion, RandomCrop, RandomCutOut,
                         RandomMosaic, RandomRotate, RandomRotFlip, Rerange,
                         ResizeShortestEdge, ResizeToMultiple, RGB2Gray,
                         SegRescale)
from .voc import PascalVOCDataset

# yapf: enable
__all__ = [
"""
    # 增加多个 vars.MyDataset_C 到 __all__
    all_datasets = ""
    for dataset_name in dataset_class_names:
        all_datasets += f"    '{dataset_name}', # TODO\n"

    # __all__ 中其他固定部分
    all_fixed = """    'BaseSegDataset', 'BioMedical3DRandomCrop', 'BioMedical3DRandomFlip',
    'CityscapesDataset', 'PascalVOCDataset', 'ADE20KDataset',
    'PascalContextDataset', 'PascalContextDataset59', 'ChaseDB1Dataset',
    'DRIVEDataset', 'HRFDataset', 'STAREDataset', 'DarkZurichDataset',
    'NightDrivingDataset', 'COCOStuffDataset', 'LoveDADataset',
    'MultiImageMixDataset', 'iSAIDDataset', 'ISPRSDataset', 'PotsdamDataset',
    'LoadAnnotations', 'RandomCrop', 'SegRescale', 'PhotoMetricDistortion',
    'RandomRotate', 'AdjustGamma', 'CLAHE', 'Rerange', 'RGB2Gray',
    'RandomCutOut', 'RandomMosaic', 'PackSegInputs', 'ResizeToMultiple',
    'LoadImageFromNDArray', 'LoadBiomedicalImageFromFile',
    'LoadBiomedicalAnnotation', 'LoadBiomedicalData', 'GenerateEdge',
    'DecathlonDataset', 'LIPDataset', 'ResizeShortestEdge',
    'BioMedicalGaussianNoise', 'BioMedicalGaussianBlur',
    'BioMedicalRandomGamma', 'BioMedical3DPad', 'RandomRotFlip',
    'SynapseDataset', 'REFUGEDataset', 'MapillaryDataset_v1',
    'MapillaryDataset_v2', 'Albu', 'LEVIRCDDataset',
    'LoadMultipleRSImageFromFile', 'LoadSingleRSImageFromFile',
    'ConcatCDInput', 'BaseCDDataset', 'DSDLSegDataset', 'BDD100KDataset',
    'NYUDataset', 'HSIDrive20Dataset'
]
"""

    # 拼接完整内容
    content = header + imports + middle + all_datasets + all_fixed

    # 写入文件
    with open(output_file, 'w', encoding='utf-8') as file:
        file.write(content)
    
    print(f"\033[93m{output_file} file generated successfully\033[0m")
    return True

if __name__ == "__main__":
    ########### 定义各参数 ###########
    # 可以定义多个数据集文件名 和 对应类名
    dataset_file_names = ["my_dataset_model"] # =['my_dataset', 'my_dataset_2'] # =["my_dataset"]
    dataset_class_names = ["MyDataset_model"] # =['MyDataset', 'MyDataset2'] # =["MyDataset"]

    ########### 文件存储位置 ###########
    output_mmseg_datasets_init = os.path.join('./mmseg/datasets/__init__.py')

    # 生成 ./mmseg/datasets/__init__.py 文件
    success = generate_mmseg_datasets_init_file(output_file=output_mmseg_datasets_init, dataset_file_names=dataset_file_names, dataset_class_names=dataset_class_names)
    

    
