import os
import torch


def save_checkpoint(path, model, optimizer=None, scaler=None, epoch=0, best_val=None, ema_state=None):
    if hasattr(model, "module"):
        model = model.module

    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)

    state = {
        "epoch": epoch,
        "model": model.state_dict(),
    }

    if optimizer:
        state["optimizer"] = optimizer.state_dict()
    if scaler:
        state["scaler"] = scaler.state_dict()
    if best_val is not None:
        state["best_val"] = best_val
    if ema_state is not None:
        state["ema"] = ema_state

    torch.save(state, path)


def load_checkpoint(path, model, optimizer=None, scaler=None):
    ckpt = torch.load(path, map_location="cpu")
    model.load_state_dict(ckpt["model"])

    if optimizer and "optimizer" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer"])
    if scaler and "scaler" in ckpt:
        scaler.load_state_dict(ckpt["scaler"])

    ema_state = ckpt.get("ema")

    return ckpt.get("epoch", 0), ckpt.get("best_val", float("inf")), ema_state
