"""优化器与学习率调度。"""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

import torch
from torch import nn


def build_optimizer(model: nn.Module, cfg: Mapping[str, Any]) -> torch.optim.Optimizer:
    """`cfg` 形如 `{name: sgd, lr: 0.01, momentum: 0.9, weight_decay: 1e-4}`。

    BatchNorm 与 bias 不做 weight decay——小数据集上这项影响明显。
    """
    cfg = dict(cfg)
    name = str(cfg.pop("name", "adamw")).lower()
    weight_decay = float(cfg.pop("weight_decay", 0.0))

    decay, no_decay = [], []
    for module in model.modules():
        for param_name, param in module.named_parameters(recurse=False):
            if not param.requires_grad:
                continue
            if param_name == "bias" or isinstance(module, nn.modules.batchnorm._BatchNorm):
                no_decay.append(param)
            else:
                decay.append(param)
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]

    if name == "sgd":
        cfg.setdefault("momentum", 0.9)
        cfg.setdefault("nesterov", True)
        return torch.optim.SGD(groups, **cfg)
    if name == "adam":
        return torch.optim.Adam(groups, **cfg)
    if name == "adamw":
        return torch.optim.AdamW(groups, **cfg)
    raise ValueError(f"未知优化器: {name}")


def build_scheduler(optimizer: torch.optim.Optimizer, cfg: Mapping[str, Any] | None, epochs: int):
    """返回按 epoch 步进的 scheduler；`cfg` 为空则不使用。"""
    if not cfg:
        return None
    cfg = dict(cfg)
    name = str(cfg.pop("name", "cosine")).lower()
    if name == "cosine":
        cfg.setdefault("T_max", epochs)
        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **cfg)
    if name == "step":
        cfg.setdefault("step_size", max(epochs // 3, 1))
        return torch.optim.lr_scheduler.StepLR(optimizer, **cfg)
    if name == "multistep":
        cfg.setdefault("milestones", [int(epochs * 0.6), int(epochs * 0.85)])
        return torch.optim.lr_scheduler.MultiStepLR(optimizer, **cfg)
    raise ValueError(f"未知调度器: {name}")