"""This module implements some workarounds for dependencies

Authors
 * Aku Rouhe 2022
"""

import warnings
import weakref

import torch

WEAKREF_MARKER = "WEAKREF"


def _cycliclrsaver(obj, path):
    state_dict = obj.state_dict()
    if state_dict.get("_scale_fn_ref") is not None:
        state_dict["_scale_fn_ref"] = WEAKREF_MARKER
    torch.save(state_dict, path)


def _cycliclrloader(obj, path, end_of_epoch):
    del end_of_epoch  # Unused
    device = "cpu"
    state_dict = torch.load(path, map_location=device)
    if state_dict.get("_scale_fn_ref") == WEAKREF_MARKER:
        if not isinstance(obj._scale_fn_ref, weakref.WeakMethod):
            MSG = "Loading CyclicLR scheduler and the _scale_ref_fn did not exist in instance."
            MSG += " You did not construct it with the same parameters it was created!"
            MSG += " Looks like you changed the scale function!"
            MSG += " If this was not intentional, the scheduler might not work correctly."
            warnings.warn(MSG)
    try:
        obj.load_state_dict(torch.load(path, map_location=device), strict=True)
    except TypeError:
        obj.load_state_dict(torch.load(path, map_location=device))