import os
import torch


def save_checkpoint(
        model,
        optimizer,
        epoch,
        loss,
        path,
):
    """
    Save a training checkpoint.
    """

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

    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "loss": loss,
        },
        path,
    )


def load_checkpoint(
        path,
        model,
        optimizer=None,
        map_location="cpu",
):
    """
    Load a training checkpoint.

    Returns:
        epoch (int)
        loss (float)
    """

    checkpoint = torch.load(
        path,
        map_location=map_location,
    )

    model.load_state_dict(
        checkpoint["model_state_dict"]
    )

    if optimizer is not None:
        optimizer.load_state_dict(
            checkpoint["optimizer_state_dict"]
        )

    return (
        checkpoint["epoch"],
        checkpoint["loss"],
    )


def latest_checkpoint(folder):
    """
    Returns the newest checkpoint file in a folder.
    """

    if not os.path.exists(folder):
        return None

    checkpoints = [
        os.path.join(folder, file)
        for file in os.listdir(folder)
        if file.endswith(".pth")
    ]

    if len(checkpoints) == 0:
        return None

    checkpoints.sort(
        key=os.path.getmtime,
        reverse=True,
    )

    return checkpoints[0]