import os

def generate_configs_base_datasets_my_dataset_file(
    output_file='./configs/_base_/datasets/my_dataset_model.py',
    dataset_class_name='MyDataset_model',
    data_root='/home/audience/Desktop/Seg_data/Data',
    img_scale=(1920, 1080),
    crop_size=(512, 512),
    train_batch_size=4,
    train_num_workers=4,
    val_and_test_batch_size=1,
    val_and_test_num_workers=4,
    train_img_path='A_Ori',
    train_seg_map_path='A_Label_GT_label_fold',
    val_img_path='A_Ori',
    val_seg_map_path='A_Label_GT_label_fold',
    test_img_path='A_Ori',
    test_seg_map_path='A_Label_GT_label_fold',
):
    # 定义模板
    dataset_config_template = f"""# dataset settings
dataset_class_name = '{dataset_class_name}'  # TODO 上一步中你定义的数据集的名字
data_root = '{data_root}'  # TODO 数据集存储路径
# img_norm_cfg = dict(
#      mean=[33.30, 35.03, 47.23], std=[48.00, 50.4, 60.51], to_rgb=True) # TODO 数据集的均值和标准差，空引用默认的，也可以网上搜代码计算
img_scale = {img_scale}  # img_scale图像尺寸 TODO (1920,1080)
crop_size = {crop_size}  # 数据增强时裁剪的大小 TODO 之后可以修改
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),  # ", reduce_zero_label=False"  TODO 是否忽略0直选项
    dict(type='RandomResize', scale=img_scale, ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    # dict(type='GenerateEdge', edge_width=4), # For pidnet
    dict(type='PackSegInputs')
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', scale=img_scale, keep_ratio=True),
    # add loading annotation after ``Resize`` because ground truth
    # does not need to do resize data transform
    dict(type='LoadAnnotations'),
    dict(type='PackSegInputs')
]
train_dataloader = dict(  # Train dataloader config
    batch_size={train_batch_size},  # Batch size of a single GPU TODO
    num_workers={train_num_workers},  # Worker to pre-fetch data for each single GPU
    persistent_workers=True,  # Shut down the worker processes after an epoch end, which can accelerate training speed.
    sampler=dict(type='DefaultSampler', shuffle=True),  # Randomly shuffle during training.
    dataset=dict(  # Train dataset config
        type=dataset_class_name,  # Type of dataset, refer to mmseg/datasets/ for details.
        data_root=data_root,  # The root of dataset.
        data_prefix=dict(
            img_path='{train_img_path}',
            seg_map_path='{train_seg_map_path}'),
        pipeline=train_pipeline))  # Processing pipeline. This is passed by the train_pipeline created before.
val_dataloader = dict(
    batch_size={val_and_test_batch_size},  # Batch size of a single GPU
    num_workers={val_and_test_num_workers},  # Worker to pre-fetch data for each single GPU
    persistent_workers=True,  # Shut down the worker processes after an epoch end, which can accelerate testing speed.
    sampler=dict(type='DefaultSampler', shuffle=False),  # Not shuffle during validation and testing.
    dataset=dict(  # Test dataset config
        type=dataset_class_name,  # Type of dataset, refer to mmseg/datasets/ for details.
        data_root=data_root,  # The root of dataset.
        data_prefix=dict(
            img_path='{val_img_path}',
            seg_map_path='{val_seg_map_path}'),
        pipeline=test_pipeline))  # Processing pipeline. This is passed by the test_pipeline created before.
test_dataloader = dict(
    batch_size={val_and_test_batch_size},  # Batch size of a single GPU
    num_workers={val_and_test_num_workers},  # Worker to pre-fetch data for each single GPU
    persistent_workers=True,  # Shut down the worker processes after an epoch end, which can accelerate testing speed.
    sampler=dict(type='DefaultSampler', shuffle=False),  # Not shuffle during validation and testing.
    dataset=dict(  # Test dataset config
        type=dataset_class_name,  # Type of dataset, refer to mmseg/datasets/ for details.
        data_root=data_root,  # The root of dataset.
        data_prefix=dict(
            img_path='{test_img_path}',
            seg_map_path='{test_seg_map_path}'),
        pipeline=test_pipeline))  # Processing pipeline. This is passed by the test_pipeline created before.
# The metric to measure the accuracy. Here, we use IoUMetric.
val_evaluator = dict(type='IoUMetric', iou_metrics=['mIoU'])
test_evaluator = val_evaluator
"""

    # 创建目录（如果不存在的话）
    os.makedirs(os.path.dirname(output_file), exist_ok=True)

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

if __name__ == '__main__':
    ########### 定义各参数 ###########
    dataset_file_name='my_dataset_model' # 数据集 文件名.py
    dataset_class_name='MyDataset_model' # 数据集 类名称
    data_root='/home/audience/Desktop/Seg_data/Data' # 数据根目录
    img_scale=(1920, 1080) # 图片大小
    # 训练、验证、测试集所在文件夹
    train_img_path='A_Ori'
    train_seg_map_path='A_Label_GT_label_fold'
    val_img_path='A_Ori'
    val_seg_map_path='A_Label_GT_label_fold'
    test_img_path='A_Ori'
    test_seg_map_path='A_Label_GT_label_fold'

    # 一般不太会变的参数
    crop_size=(512, 512) # 分割大小
    train_batch_size=4 # 训练batch
    train_num_workers=4 # 训练并行运行数量
    val_and_test_batch_size=1 # 验证集和测试集batch
    val_and_test_num_workers=4 # 验证集和测试集并行运行数量
    
    ########### 文件存储位置 ###########
    output_configs_base_datasets_my_dataset=f'./configs/_base_/datasets/{dataset_file_name}.py'

    # 使用默认变量生成配置文件
    success = generate_configs_base_datasets_my_dataset_file(output_file=output_configs_base_datasets_my_dataset, dataset_class_name=dataset_class_name , data_root=data_root, img_scale=img_scale, crop_size=crop_size, train_batch_size=train_batch_size, train_num_workers=train_num_workers, val_and_test_batch_size=val_and_test_batch_size, val_and_test_num_workers=val_and_test_num_workers, train_img_path=train_img_path, train_seg_map_path=train_seg_map_path, val_img_path=val_img_path, val_seg_map_path=val_seg_map_path, test_img_path=test_img_path, test_seg_map_path=test_seg_map_path)
