import os

def generate_epochs_configs_base_schedules_schedule_file(output_file, max_epochs=300, val_interval=1, checkpoint_interval=10, loggerhook_interval=300):
    """
    生成基于 Epoch 的 MMSegmentation schedule 配置文件。

    Args:
        output_file (str): 输出配置文件的路径。
        max_epochs (int): 最大训练轮次。
        val_interval (int): 验证间隔的轮次数。
        checkpoint_interval (int): 保存模型权重间隔的轮次数。
        loggerhook_interval (int): 日志打印的迭代间隔。
    """
    
    # 定义 Epoch-based 的模板
    schedule_config_template = f"""# optimizer
# For SGD.
# optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
# For AdamW.
optimizer = dict(type='AdamW', lr=0.0001, betas=(0.9, 0.999), weight_decay=0.01)
optim_wrapper = dict(type='OptimWrapper', optimizer=optimizer, clip_grad=None)

# learning policy
param_scheduler = [
    dict(
        type='PolyLR',
        eta_min=1e-4,
        power=0.9,
        begin=0,
        end={max_epochs},
        by_epoch=True)
]

# training schedule for {max_epochs} epochs
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs={max_epochs}, val_interval={val_interval})
val_cfg = dict(type='ValLoop')
test_cfg = dict(type='TestLoop')

default_hooks = dict(
    timer=dict(type='IterTimerHook'),
    logger=dict(type='LoggerHook', interval={loggerhook_interval}, log_metric_by_epoch=True),
    param_scheduler=dict(type='ParamSchedulerHook'),
    checkpoint=dict(type='CheckpointHook', by_epoch=True, interval={checkpoint_interval}, save_best='mIoU', rule='greater'),
    sampler_seed=dict(type='DistSamplerSeedHook'),
    visualization=dict(type='SegVisualizationHook'))
"""

    # --- 将生成的内容输出到界面 ---
    print("\033[92m--- Generated Configuration Content ---\033[0m")
    print(schedule_config_template)
    print("\033[92m-------------------------------------\033[0m")

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

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

if __name__ == '__main__':
    ########### 定义各参数 ###########
    max_epochs = 300          # 训练总轮数
    val_interval = 1          # 每 1 个 epoch 验证一次
    checkpoint_interval = 10  # 每 10 个 epoch 保存一次模型
    loggerhook_interval = max_epochs   # 日志间隔1次迭代

    ########### 文件存储位置 ###########
    # 文件名将反映 epoch 数量, e.g., schedule_300e.py
    output_filename = f'schedule_{max_epochs}e.py'
    output_configs_base_schedules_file = os.path.join('./configs/_base_/schedules/', output_filename)

    # 使用变量生成配置文件
    success = generate_epoch_based_schedule_file(
        output_file=output_configs_base_schedules_file,
        max_epochs=max_epochs,
        val_interval=val_interval,
        checkpoint_interval=checkpoint_interval,
        loggerhook_interval=loggerhook_interval
    )