# Copyright (c) OpenMMLab. All rights reserved.
from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
                            LoggerHook, ParamSchedulerHook)
from mmengine.optim import AmpOptimWrapper # 导入 AmpOptimWrapper
from mmengine.optim.optimizer.optimizer_wrapper import OptimWrapper
from mmengine.optim.scheduler.lr_scheduler import PolyLR
# 导入 EpochBasedTrainLoop
from mmengine.runner.loops import EpochBasedTrainLoop, TestLoop, ValLoop
from torch.optim.adamw import AdamW # 推荐使用 AdamW 优化器

from mmseg.engine import SegVisualizationHook

# --- 修改部分 ---

# 1. 优化器 (Optimizer)
# 推荐使用 AdamW，它通常比 SGD 效果更好且更稳定
optimizer = dict(
    type=AdamW,
    lr=0.0001, # AdamW 的学习率通常设置得比 SGD 小
    betas=(0.9, 0.999),
    weight_decay=0.01)

# 优化器封装 (OptimWrapper)
# clip_grad 用于梯度裁剪，防止梯度爆炸，可以根据需要设置
optim_wrapper = dict(
    type=OptimWrapper, optimizer=optimizer, clip_grad=dict(max_norm=1, norm_type=2))

# 2. 学习率调度器 (Learning Rate Scheduler)
# 总轮次设为 200 epochs
max_epochs = 200
param_scheduler = [
    dict(
        type=PolyLR,
        eta_min=1e-5, # 学习率最小值
        power=0.9,
        begin=0,
        end=max_epochs, # 关键修改：结束点改为总 epochs
        by_epoch=True) # 关键修改：改为按 epoch 更新学习率
]

# 3. 训练、验证和测试的配置
# 关键修改：使用 EpochBasedTrainLoop
train_cfg = dict(type=EpochBasedTrainLoop, max_epochs=max_epochs, val_interval=1) # 每 1 个 epoch 验证一次
val_cfg = dict(type=ValLoop)
test_cfg = dict(type=TestLoop)

# 4. 默认钩子 (Default Hooks)
default_hooks = dict(
    timer=dict(type=IterTimerHook),
    # 关键修改：让日志按 epoch 记录
    logger=dict(type=LoggerHook, interval=300, log_metric_by_epoch=True),
    param_scheduler=dict(type=ParamSchedulerHook),
    # 关键修改：让检查点（模型权重）按 epoch 保存
    checkpoint=dict(type=CheckpointHook, by_epoch=True, interval=10), # 每 10 个 epoch 保存一次
    sampler_seed=dict(type=DistSamplerSeedHook),
    visualization=dict(type=SegVisualizationHook))