from __future__ import annotations

from pathlib import Path

import torch


def save(path: Path, model, optimizer, scheduler, step: int, best: float, extra: dict | None = None) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict() if optimizer else None,
                "scheduler": scheduler.state_dict() if scheduler else None, "step": step, "best": best,
                "cfg": model.cfg.__dict__, "extra": extra or {}}, tmp)
    tmp.replace(path)


def load(path: Path, model, optimizer=None, scheduler=None, map_location="cpu") -> dict:
    ck = torch.load(path, map_location=map_location, weights_only=False)
    model.load_state_dict(ck["model"])
    if optimizer and ck.get("optimizer"):
        optimizer.load_state_dict(ck["optimizer"])
    if scheduler and ck.get("scheduler"):
        scheduler.load_state_dict(ck["scheduler"])
    return ck