from abc import ABC
from abc import abstractmethod

import torch.optim.lr_scheduler as L


class AbsScheduler(ABC):
    @abstractmethod
    def step(self, epoch: int = None):
        pass

    @abstractmethod
    def state_dict(self):
        pass

    @abstractmethod
    def load_state_dict(self, state):
        pass


# If you need to define custom scheduler, please inherit these classes
class AbsBatchStepScheduler(AbsScheduler):
    @abstractmethod
    def step(self, epoch: int = None):
        pass

    @abstractmethod
    def state_dict(self):
        pass

    @abstractmethod
    def load_state_dict(self, state):
        pass


class AbsEpochStepScheduler(AbsScheduler):
    @abstractmethod
    def step(self, epoch: int = None):
        pass

    @abstractmethod
    def state_dict(self):
        pass

    @abstractmethod
    def load_state_dict(self, state):
        pass


class AbsValEpochStepScheduler(AbsEpochStepScheduler):
    @abstractmethod
    def step(self, val, epoch: int = None):
        pass

    @abstractmethod
    def state_dict(self):
        pass

    @abstractmethod
    def load_state_dict(self, state):
        pass


# Create alias type to check the type
# Note(kamo): Currently PyTorch doesn't provide the base class
# to judge these classes.
AbsValEpochStepScheduler.register(L.ReduceLROnPlateau)
for s in [
    L.ReduceLROnPlateau,
    L.LambdaLR,
    L.StepLR,
    L.MultiStepLR,
    L.MultiStepLR,
    L.ExponentialLR,
    L.CosineAnnealingLR,
]:
    AbsEpochStepScheduler.register(s)

AbsBatchStepScheduler.register(L.CyclicLR)
for s in [
    L.OneCycleLR,
    L.CosineAnnealingWarmRestarts,
]:
    AbsBatchStepScheduler.register(s)