"""检查点存取。"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
from torch import nn


def save_checkpoint(
    path: str | Path,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
    epoch: int = 0,
    best_metric: float | None = None,
    config: dict | None = None,
    scheduler: Any = None,
    global_step: int = 0,
    train_cfg: dict | None = None,
) -> Path:
    """保存检查点。

    除模型与优化器外还记下调度器状态、global_step 和当时的 train 配置：
    少了这些，续训时学习率调度会从头开始，而且无法判断 yaml 里的超参有没有被改过。
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict() if optimizer else None,
            "scheduler": scheduler.state_dict() if scheduler else None,
            "epoch": epoch,
            "global_step": global_step,
            "best_metric": best_metric,
            "config": config,
            "train_cfg": train_cfg,
        },
        path,
    )
    return path


def load_checkpoint(
    path: str | Path,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
    map_location: Any = "cpu",
    strict: bool = True,
) -> dict:
    state = torch.load(Path(path), map_location=map_location, weights_only=False)
    model.load_state_dict(state["model"], strict=strict)
    if optimizer is not None and state.get("optimizer"):
        optimizer.load_state_dict(state["optimizer"])
    return state