import os

def generate_times_configs_base_schedules_schedule_file(output_file, train_time_k=4, val_proportion=1/10, loggerhook_interval=50):
    
    # 定义模板
    schedule_config_template = f"""# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
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={int(train_time_k*1000)},
        by_epoch=False)
]
# training schedule for {train_time_k}k
train_cfg = dict(
    type='IterBasedTrainLoop', max_iters={int(train_time_k*1000)}, val_interval={int(train_time_k*1000*val_proportion)})
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=False),
    param_scheduler=dict(type='ParamSchedulerHook'),
    checkpoint=dict(type='CheckpointHook', by_epoch=False, interval={int(train_time_k*1000*val_proportion)}, , save_best='mIoU', rule='greater'),
    sampler_seed=dict(type='DistSamplerSeedHook'),
    visualization=dict(type='SegVisualizationHook'))
"""

    # 创建目录（如果不存在的话）
    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__':
    ########### 定义各参数 ###########
    train_time_k = 4  # 训练轮数（以k为单位）

    # 一般不太会变的参数
    val_proportion = 1/10 # 验证比例
    loggerhook_interval = 50  # 日志间隔50次迭代

    checkpoint_interval = int(train_time_k*1000*val_proportion) # 训练多少轮保存pth # TODO 可自定义

    ########### 文件存储位置 ###########
    output_configs_base_schedules_schedules_Timek = os.path.join('./configs/_base_/schedules/', f'schedule_{train_time_k}k.py')

    # 使用变量生成配置文件
    success = generate_configs_base_schedules_schedule_file(output_file=output_configs_base_schedules_schedules_Timek, train_time_k=train_time_k, val_proportion=val_proportion, loggerhook_interval=loggerhook_interval)
