from pathlib import Path
import torch, time, shutil

def save_checkpoint(path, model, optimizer=None, scheduler=None, step=0, epoch=0, metrics=None, config=None):
    path=Path(path); path.parent.mkdir(parents=True, exist_ok=True)
    tmp=path.with_suffix(path.suffix+'.tmp')
    state={'model':model.state_dict(),'step':step,'epoch':epoch,'metrics':metrics or {},'saved_at':time.time(),'config':config}
    if optimizer is not None: state['optimizer']=optimizer.state_dict()
    if scheduler is not None: state['scheduler']=scheduler.state_dict()
    torch.save(state,tmp); tmp.replace(path)
    latest=path.parent/'latest.pt'
    if latest != path: shutil.copy2(path, latest)

def load_checkpoint(path, model, optimizer=None, scheduler=None, map_location='cpu'):
    state=torch.load(path,map_location=map_location)
    model.load_state_dict(state['model'])
    if optimizer is not None and 'optimizer' in state: optimizer.load_state_dict(state['optimizer'])
    if scheduler is not None and 'scheduler' in state: scheduler.load_state_dict(state['scheduler'])
    return state