import torch

from allennlp.common.params import Params
from allennlp.training.callbacks.callback import Callback, handle_event
from allennlp.training.callbacks.events import Events
from allennlp.training.learning_rate_schedulers import LearningRateScheduler


@Callback.register("update_learning_rate")
class UpdateLearningRate(Callback):
    """
    Callback that runs the learning rate scheduler.

    # Parameters

    learning_rate_scheduler : `LearningRateScheduler`
        The scheduler to handler.
    """

    def __init__(self, learning_rate_scheduler: LearningRateScheduler) -> None:
        self.learning_rate_scheduler = learning_rate_scheduler

    @handle_event(Events.BACKWARD, priority=1000)
    def step_batch(self, trainer):
        self.learning_rate_scheduler.step_batch(trainer.batch_num_total)

    @handle_event(Events.EPOCH_END)
    def step(self, trainer):
        self.learning_rate_scheduler.step(trainer.latest_val_metric, trainer.epoch_number)

    def get_training_state(self) -> dict:
        """
        We need to persist the learning_rate_scheduler state as training state.
        """
        return {"learning_rate_scheduler": self.learning_rate_scheduler.state_dict()}

    def restore_training_state(self, training_state: dict) -> None:
        state_dict = training_state.pop("learning_rate_scheduler", None)

        if state_dict:
            self.learning_rate_scheduler.load_state_dict(state_dict)

    @classmethod
    def from_params(  # type: ignore
        cls, params: Params, optimizer: torch.optim.Optimizer, **extras
    ) -> "UpdateLearningRate":

        return cls(
            LearningRateScheduler.from_params(
                params=params.pop("learning_rate_scheduler"), optimizer=optimizer
            )
        )
