# coding=utf-8
__author__ = "Dimitris Karkalousos"

# Taken and adapted from: https://github.com/NVIDIA/NeMo/blob/main/nemo/core/optim/lr_scheduler.py

import copy
import dataclasses
import inspect
import math
import warnings
from functools import partial
from typing import Any, Dict, Optional, Union

import hydra
import torch.optim.lr_scheduler as pt_scheduler
from omegaconf import DictConfig, OmegaConf
from torch import optim
from torch.optim.lr_scheduler import _LRScheduler
from torch.utils.data import dataloader

from atommic.core.conf.schedulers import SchedulerParams, get_scheduler_config, register_scheduler_params
from atommic.utils import logging
from atommic.utils.model_utils import maybe_update_config_version


class WarmupPolicy(_LRScheduler):
    """Adds warmup kwargs and warmup logic to lr policy. All arguments should be passed as kwargs for clarity.

    Returns
    -------
    lr : float
        Learning rate for current step.
    """

    def __init__(self, optimizer, *, warmup_steps=None, warmup_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1):
        """Inits :class:`WarmupPolicy`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        warmup_steps : int
            Number of training steps in warmup stage. Default is ``None``.
        warmup_ratio : float
            Ratio of warmup steps to total steps. Default is ``None``.
        max_steps : int
            Total number of steps while training or `None` for infinite training. Default is ``None``.
        min_lr : float
            Minimum learning rate. Default is ``0``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        if warmup_steps is not None and warmup_ratio is not None:
            raise AssertionError("Either use particular number of step or ratio")
        if warmup_ratio is not None and max_steps is None:
            raise AssertionError("If there is a ratio, there should be a total steps")

        # It is necessary to assign all attributes *before* __init__,
        # as class is wrapped by an inner class.
        self.max_steps = max_steps
        if warmup_steps is not None:
            self.warmup_steps = warmup_steps
        elif warmup_ratio is not None:
            self.warmup_steps = int(warmup_ratio * max_steps)
        else:
            self.warmup_steps = 0

        self.min_lr = min_lr
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        """Get learning rate at current step."""
        if not self._get_lr_called_within_step:
            warnings.warn(
                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
            )

        step = self.last_epoch

        if step <= self.warmup_steps and self.warmup_steps > 0:
            return self._get_warmup_lr(step)

        if (self.max_steps is not None) and (step > self.max_steps):
            return [self.min_lr for _ in self.base_lrs]

        return self._get_lr(step)

    def _get_warmup_lr(self, step):
        """Linear warmup"""
        lr_val = (step + 1) / (self.warmup_steps + 1)
        return [initial_lr * lr_val for initial_lr in self.base_lrs]

    def _get_lr(self, step):  # pylint: disable=unused-argument
        """Simple const lr policy"""
        return self.base_lrs


class SquareRootConstantPolicy(_LRScheduler):
    """Adds warmup kwargs and warmup logic to lr policy. All arguments should be passed as kwargs for clarity."""

    def __init__(
        self, optimizer, *, constant_steps=None, constant_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1
    ):
        """Inits :class:`SquareRootConstantPolicy`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        constant_steps : int
            Number of training steps in constant stage. Default is ``None``.
        constant_ratio : float
            Ratio of constant steps to total steps. Default is ``None``.
        max_steps : int
            Total number of steps while training or `None` for infinite training. Default is ``None``.
        min_lr : float
            Minimum learning rate. Default is ``0``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        if constant_steps is not None and constant_ratio is not None:
            raise AssertionError("Either use particular number of step or ratio")

        if constant_ratio is not None and max_steps is None:
            raise AssertionError("If there is a ratio, there should be a total steps")

        # It is necessary to assign all attributes *before* __init__, as class is wrapped by an inner class.
        self.max_steps = max_steps
        if constant_steps is not None:
            self.constant_steps = constant_steps
        elif constant_ratio is not None:
            self.constant_steps = int(constant_ratio * max_steps)
        else:
            self.constant_steps = 0

        self.constant_lr = 1 / (constant_steps**0.5)
        self.min_lr = min_lr
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        """Get learning rate at current step."""
        if not self._get_lr_called_within_step:
            warnings.warn(
                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
            )

        step = self.last_epoch

        if step <= self.constant_steps:
            return [self.constant_lr for _ in self.base_lrs]

        if step > self.max_steps:
            return [self.min_lr for _ in self.base_lrs]

        return self._get_lr(step)

    def _get_lr(self, step):  # pylint: disable=unused-argument
        """Simple const lr policy"""
        return self.base_lrs


class WarmupHoldPolicy(WarmupPolicy):
    """Variant of WarmupPolicy which maintains high learning rate for a defined number of steps. All arguments should
    be passed as kwargs for clarity,

    Results
    -------
    lr : float
        Learning rate is linearly increased from 0 to 1 over warmup steps, then linearly decreased from 1 to 0 over
        hold steps.
    """

    def __init__(
        self,
        optimizer,
        *,
        warmup_steps=None,
        warmup_ratio=None,
        hold_steps=None,
        hold_ratio=None,
        max_steps=None,
        min_lr=0.0,
        last_epoch=-1,
    ):
        """Inits :class:`WarmupHoldPolicy`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        warmup_steps : int
            Number of training steps in warmup stage. Default is ``None``.
        warmup_ratio : float
            Ratio of warmup steps to total steps. Default is ``None``.
        hold_steps : int
            Number of training steps to hold the learning rate after warm up. Default is ``None``.
        hold_ratio : float
            Ratio of hold steps to total steps. Default is ``None``.
        max_steps : int
            Total number of steps while training or `None` for infinite training. Default is ``None``.
        min_lr : float
            Minimum learning rate. Default is ``0``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        if hold_steps is not None and hold_ratio is not None:
            raise AssertionError("Either use particular number of step or ratio")
        if hold_ratio is not None and max_steps is None:
            raise AssertionError("If there is a ratio, there should be a total steps")

        self.min_lr = min_lr
        self._last_warmup_lr = 0.0

        # Necessary to duplicate as class attributes are hidden in inner class
        self.max_steps = max_steps
        if warmup_steps is not None:
            self.warmup_steps = warmup_steps
        elif warmup_ratio is not None:
            self.warmup_steps = int(warmup_ratio * max_steps)
        else:
            self.warmup_steps = 0

        if hold_steps is not None:
            self.hold_steps = hold_steps + self.warmup_steps
        elif hold_ratio is not None:
            self.hold_steps = int(hold_ratio * max_steps) + self.warmup_steps
        else:
            self.hold_steps = 0

        super().__init__(
            optimizer,
            warmup_steps=warmup_steps,
            warmup_ratio=warmup_ratio,
            max_steps=max_steps,
            last_epoch=last_epoch,
            min_lr=min_lr,
        )

    def get_lr(self):
        """Get learning rate at current step."""
        if not self._get_lr_called_within_step:
            warnings.warn(
                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
            )

        step = self.last_epoch

        # Warmup phase
        if 0 < self.warmup_steps >= step:
            return self._get_warmup_lr(step)

        # Hold phase
        if self.hold_steps < step >= self.warmup_steps:
            return self.base_lrs

        if step > self.max_steps:
            return [self.min_lr for _ in self.base_lrs]

        return self._get_lr(step)


class WarmupAnnealHoldPolicy(_LRScheduler):
    """Adds warmup kwargs and warmup logic to lr policy. All arguments should be passed as kwargs for clarity."""

    def __init__(
        self,
        optimizer,
        *,
        warmup_steps=None,
        warmup_ratio=None,
        constant_steps=None,
        constant_ratio=None,
        max_steps=None,
        min_lr=0.0,
        last_epoch=-1,
    ):
        """Inits :class:`WarmupAnnealHoldPolicy`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        warmup_steps : int
            Number of training steps in warmup stage. Default is ``None``.
        warmup_ratio : float
            Ratio of warmup steps to total steps. Default is ``None``.
        constant_steps : int
            Number of training steps in constant stage. Default is ``None``.
        constant_ratio : float
            Ratio of constant steps to total steps. Default is ``None``.
        max_steps : int
            Total number of steps while training or `None` for infinite training. Default is ``None``.
        min_lr : float
            Minimum learning rate. Default is ``0``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        if warmup_steps is not None and warmup_ratio is not None:
            raise AssertionError("Either use particular number of step or ratio")
        if constant_steps is not None and constant_ratio is not None:
            raise AssertionError("Either use constant_steps or constant_ratio")
        if warmup_ratio is not None and max_steps is None:
            raise AssertionError("If there is a ratio, there should be a total steps")

        # It is necessary to assign all attributes *before* __init__, as class is wrapped by an inner class.
        self.max_steps = max_steps

        if warmup_steps is not None:
            self.warmup_steps = warmup_steps
        elif warmup_ratio is not None:
            self.warmup_steps = int(warmup_ratio * max_steps)
        else:
            self.warmup_steps = 0

        if constant_steps is not None:
            self.constant_steps = constant_steps
        elif constant_ratio is not None:
            self.constant_steps = int(constant_ratio * max_steps)
        else:
            self.constant_steps = 0

        self.decay_steps = max_steps - (self.constant_steps + self.warmup_steps)

        self.min_lr = min_lr
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        """Get learning rate at current step."""
        if not self._get_lr_called_within_step:
            warnings.warn(
                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
            )

        step = self.last_epoch

        # Warmup steps
        if 0 < self.warmup_steps >= step:
            return self._get_warmup_lr(step)

        # Constant steps after warmup and decay
        if self.constant_steps > 0 and (self.warmup_steps + self.decay_steps) < step <= self.max_steps:
            return self._get_constant_lr(step)

        # Min lr after max steps of updates
        if step > self.max_steps:
            return [self.min_lr for _ in self.base_lrs]

        return self._get_lr(step)

    def _get_warmup_lr(self, step):
        """Get learning rate at warmup stage."""
        lr_val = (step + 1) / (self.warmup_steps + 1)
        return [initial_lr * lr_val for initial_lr in self.base_lrs]

    def _get_constant_lr(self, step):  # pylint: disable=unused-argument
        """Get learning rate at constant stage."""
        return [self.min_lr for _ in self.base_lrs]

    def _get_lr(self, step):  # pylint: disable=unused-argument
        """Simple const lr policy"""
        return self.base_lrs


def _sqrt_annealing(initial_lr, step, max_steps, min_lr):
    """Anneal learning rate by sqrt."""
    mult = ((max_steps - step) / max_steps) ** 0.5
    out_lr = initial_lr * mult
    out_lr = max(out_lr, min_lr)
    return out_lr


def _square_annealing(initial_lr, step, max_steps, min_lr):
    """Anneal learning rate by square."""
    mult = ((max_steps - step) / max_steps) ** 2
    out_lr = initial_lr * mult
    out_lr = max(out_lr, min_lr)
    return out_lr


def _cosine_annealing(initial_lr, step, max_steps, min_lr):
    """Anneal learning rate by cosine."""
    mult = 0.5 * (1 + math.cos(math.pi * step / max_steps))
    return (initial_lr - min_lr) * mult + min_lr


def _linear_warmup_with_cosine_annealing(max_lr, warmup_steps, step, decay_steps, min_lr):
    """Anneal learning rate by linear warmup and cosine annealing."""
    if max_lr <= min_lr:
        raise AssertionError
    # Use linear warmup for the initial part.
    if warmup_steps > 0 and step <= warmup_steps:
        return max_lr * float(step) / float(warmup_steps)

    # For any steps larger than `decay_steps`, use `min_lr`.
    if step > warmup_steps + decay_steps:
        return min_lr

    # If we are done with the warmup period, use the decay style.
    num_steps_ = step - warmup_steps
    decay_steps_ = decay_steps
    decay_ratio = float(num_steps_) / float(decay_steps_)
    if decay_ratio < 0.0:
        raise AssertionError
    if decay_ratio > 1.0:
        raise AssertionError
    delta_lr = max_lr - min_lr

    coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0)

    return min_lr + coeff * delta_lr


def _poly_decay(initial_lr, step, decay_steps, power, min_lr, cycle):
    """Polynomial decay of learning rate."""
    if cycle:
        multiplier = 1.0 if step == 0 else math.ceil(step / decay_steps)
        decay_steps = decay_steps * multiplier
    else:
        step = min(step, decay_steps)
    p = step / decay_steps
    lr = (initial_lr - min_lr) * math.pow(1.0 - p, power)
    lr += min_lr
    return lr


def _noam_hold_annealing(initial_lr, step, warmup_steps, hold_steps, decay_rate, min_lr):
    """Anneal learning rate by noam hold."""
    # hold_steps = total number of steps to hold the LR, not the warmup + hold steps.
    T_warmup_decay = max(1, warmup_steps**decay_rate)
    T_hold_decay = max(1, (step - hold_steps) ** decay_rate)
    lr = (initial_lr * T_warmup_decay) / T_hold_decay
    return max(lr, min_lr)


class SquareAnnealing(WarmupPolicy):
    """Anneal learning rate by square."""

    def __init__(self, optimizer, *, max_steps, min_lr=1e-5, last_epoch=-1, **kwargs):
        """Inits :class:`SquareAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        max_steps : int
            Total number of steps while training or `None` for infinite training.
        min_lr : float
            Minimum learning rate. Default is ``1e-5``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        return [
            _square_annealing(
                initial_lr=initial_lr,
                step=step - self.warmup_steps,
                max_steps=self.max_steps - self.warmup_steps,
                min_lr=self.min_lr,
            )
            for initial_lr in self.base_lrs
        ]


class SquareRootAnnealing(WarmupPolicy):
    """Anneal learning rate by square root."""

    def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs):
        """Inits :class:`SquareRootAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        max_steps : int
            Total number of steps while training or `None` for infinite training.
        min_lr : float
            Minimum learning rate. Default is ``0``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        return [
            _sqrt_annealing(
                initial_lr=initial_lr,
                step=step,
                max_steps=self.max_steps,
                min_lr=self.min_lr,
            )
            for initial_lr in self.base_lrs
        ]


class CosineAnnealing(WarmupAnnealHoldPolicy):
    """Anneal learning rate by cosine."""

    def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs):
        """Inits :class:`CosineAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        max_steps : int
            Total number of steps while training or `None` for infinite training.
        min_lr : float
            Minimum learning rate. Default is ``0``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        for initial_lr in self.base_lrs:
            if initial_lr < self.min_lr:
                raise ValueError(
                    f"{self} received an initial learning rate that was lower than the minimum learning rate."
                )

        return (
            [
                _cosine_annealing(
                    initial_lr=initial_lr,
                    step=step - self.warmup_steps,
                    max_steps=self.max_steps - self.warmup_steps,
                    min_lr=self.min_lr,
                )
                for initial_lr in self.base_lrs
            ]
            if self.constant_steps is None or self.constant_steps == 0
            else self._get_linear_warmup_with_cosine_annealing_lr(step)
        )

    def _get_warmup_lr(self, step):
        """Get the warmup learning rate for the given step."""
        if self.constant_steps is None or self.constant_steps == 0:
            return super()._get_warmup_lr(step)

        # Use linear warmup for the initial part.
        return self._get_linear_warmup_with_cosine_annealing_lr(step)

    def _get_constant_lr(self, step):
        """Only called when constant_steps is not None and not 0."""
        return self._get_linear_warmup_with_cosine_annealing_lr(step)

    def _get_linear_warmup_with_cosine_annealing_lr(self, step):
        """Cosine Schedule, slightly different warmup schedule + constant LR at the end."""
        return [
            _linear_warmup_with_cosine_annealing(
                max_lr=self.base_lrs[0],
                warmup_steps=self.warmup_steps,
                step=step,
                decay_steps=self.decay_steps,
                min_lr=self.min_lr,
            )
            for _ in self.base_lrs
        ]


class NoamAnnealing(_LRScheduler):
    """Noam learning rate annealing."""

    def __init__(
        self, optimizer, *, d_model, warmup_steps=None, warmup_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1
    ):
        """Inits :class:`NoamAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer.
        d_model : int
            Model dimensionality.
        warmup_steps : int
            Number of training steps in warmup stage. Default is ``None``.
        warmup_ratio : float
            Ratio of warmup steps to total steps. Default is ``None``.
        max_steps : int
            Total number of steps while training or `None` for infinite training.
        min_lr : float
            Minimum learning rate. Default is ``0``.
        last_epoch : int
            Last epoch. Default is ``-1``.
        """
        self._normalize = d_model ** (-0.5)
        if warmup_steps is not None and warmup_ratio is not None:
            raise AssertionError("Either use particular number of step or ratio")
        if warmup_ratio is not None and max_steps is None:
            raise AssertionError("If there is a ratio, there should be a total steps")

        # It is necessary to assign all attributes *before* __init__,
        # as class is wrapped by an inner class.
        self.max_steps = max_steps
        if warmup_steps is not None:
            self.warmup_steps = warmup_steps
        elif warmup_ratio is not None:
            self.warmup_steps = int(warmup_ratio * max_steps)
        else:
            self.warmup_steps = 0

        self.min_lr = min_lr
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        """Get learning rate at current step."""
        if not self._get_lr_called_within_step:
            warnings.warn(
                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
            )

        step = max(1, self.last_epoch)

        if step > self.max_steps:
            return [self.min_lr for _ in self.base_lrs]

        for initial_lr in self.base_lrs:
            if initial_lr < self.min_lr:
                raise ValueError(
                    f"{self} received an initial learning rate that was lower than the minimum learning rate."
                )

        return [self._noam_annealing(initial_lr=initial_lr, step=step) for initial_lr in self.base_lrs]

    def _noam_annealing(self, initial_lr, step):
        """Noam learning rate annealing."""
        mult = (
            self._normalize * min(step ** (-0.5), step * (self.warmup_steps ** (-1.5)))
            if self.warmup_steps > 0
            else self._normalize * step ** (-0.5)
        )
        out_lr = initial_lr * mult
        if step > self.warmup_steps:
            out_lr = max(out_lr, self.min_lr)
        return out_lr


class NoamHoldAnnealing(WarmupHoldPolicy):
    """Implementation of the Noam Hold Annealing policy from the SqueezeFormer paper.

    Unlike NoamAnnealing, the peak learning rate can be explicitly set for this scheduler.
    The schedule first performs linear warmup, then holds the peak LR, then decays with some schedule for
    the remainder of the steps. Therefore, the min-lr is still dependent on the hyperparameters selected.

    It's schedule is determined by three factors-

    Warmup Steps: Initial stage, where linear warmup occurs uptil the peak LR is reached. Unlike NoamAnnealing,
        the peak LR is explicitly stated here instead of a scaling factor.

    Hold Steps: Intermediate stage, where the peak LR is maintained for some number of steps. In this region,
        the high peak LR allows the model to converge faster if training is stable. However the high LR
        may also cause instability during training. Should usually be a significant fraction of training
        steps (around 30-40% of the entire training steps).

    Decay Steps: Final stage, where the LR rapidly decays with some scaling rate (set by decay rate).
        To attain Noam decay, use 0.5, for Squeezeformer recommended decay, use 1.0. The fast decay after
        prolonged high LR during hold phase allows for rapid convergence.

    References
    ----------
        [1]
        [Squeezeformer: An Efficient Transformer for Automatic Speech Recognition](https://arxiv.org/abs/2206.00888)
    """

    def __init__(self, optimizer, *, max_steps, decay_rate=0.5, min_lr=0.0, last_epoch=-1, **kwargs):
        """Inits :class:`NoamHoldAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer to use for the scheduler.
        max_steps : int
            Total number of training steps.
        decay_rate : float
            Decay rate for the final stage of the schedule. Should be between 0 and 1. Default is ``0.5``.
        min_lr : float
            Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
        last_epoch : int
            Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
        """
        self.decay_rate = decay_rate
        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)

    def _get_lr(self, step):
        """Get the learning rate for the given step."""
        if self.warmup_steps is None or self.warmup_steps == 0:
            raise ValueError("Noam scheduler cannot be used without warmup steps")

        if self.hold_steps > 0:
            hold_steps = self.hold_steps - self.warmup_steps
        else:
            hold_steps = 0

        return [
            _noam_hold_annealing(
                initial_lr,
                step=step,
                warmup_steps=self.warmup_steps,
                hold_steps=hold_steps,
                decay_rate=self.decay_rate,
                min_lr=self.min_lr,
            )
            for initial_lr in self.base_lrs
        ]


class WarmupAnnealing(WarmupPolicy):
    """Warmup learning rate annealing."""

    def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
        """Inits :class:`WarmupAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer to use for the scheduler.
        max_steps : int
            Total number of training steps.
        last_epoch : int
            Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
        min_lr : float
            Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
        """
        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        delta_lr = self.base_lrs[0] - self.min_lr
        mult = (step - self.warmup_steps) / (self.max_steps - self.warmup_steps)
        return [self.min_lr + (1 - mult) * delta_lr for _ in self.base_lrs]


class InverseSquareRootAnnealing(WarmupPolicy):
    """Inverse square root learning rate annealing."""

    def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
        """Inits :class:`InverseSquareRootAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer to use for the scheduler.
        max_steps : int
            Total number of training steps.
        last_epoch : int
            Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
        min_lr : float
            Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
        """
        super().__init__(optimizer=optimizer, max_steps=max_steps, **kwargs, last_epoch=last_epoch, min_lr=min_lr)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        denom = ((step + 1) / (self.warmup_steps + 1)) ** 0.5
        return [initial_lr / denom for initial_lr in self.base_lrs]


class T5InverseSquareRootAnnealing(SquareRootConstantPolicy):
    """Inverse square root learning rate annealing."""

    def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
        """Inits :class:`T5InverseSquareRootAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer to use for the scheduler.
        max_steps : int
            Total number of training steps.
        last_epoch : int
            Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
        min_lr : float
            Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
        """
        super().__init__(optimizer=optimizer, max_steps=max_steps, **kwargs, last_epoch=last_epoch, min_lr=min_lr)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        return [1 / (step**0.5) for _ in self.base_lrs]


class PolynomialDecayAnnealing(WarmupPolicy):
    """Polynomial decay learning rate annealing."""

    def __init__(self, optimizer, *, max_steps, min_lr=0.0, power=1.0, cycle=False, last_epoch=-1, **kwargs):
        """Inits :class:`PolynomialDecayAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer to use for the scheduler.
        max_steps : int
            Total number of training steps.
        min_lr : float
            Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
        power : float
            Power of the polynomial. Default is ``1.0``.
        cycle : bool
            Whether to cycle the schedule. Default is ``False``.
        last_epoch : int
            Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
        """
        self.power = power
        self.cycle = cycle

        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        return [
            _poly_decay(
                initial_lr,
                step=step - self.warmup_steps,
                decay_steps=self.max_steps - self.warmup_steps,
                power=self.power,
                min_lr=self.min_lr,
                cycle=self.cycle,
            )
            for initial_lr in self.base_lrs
        ]


class PolynomialHoldDecayAnnealing(WarmupHoldPolicy):
    """Polynomial decay learning rate annealing."""

    def __init__(self, optimizer, *, max_steps, min_lr=0.0, power=1.0, cycle=False, last_epoch=-1, **kwargs):
        """Inits :class:`PolynomialHoldDecayAnnealing`.

        Parameters
        ----------
        optimizer : torch.optim.Optimizer
            Optimizer to use for the scheduler.
        max_steps : int
            Total number of training steps.
        min_lr : float
            Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
        power : float
            Power of the polynomial. Default is ``1.0``.
        cycle : bool
            Whether to cycle the schedule. Default is ``False``.
        last_epoch : int
            Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
        """
        self.power = power
        self.cycle = cycle

        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)

    def _get_lr(self, step):
        """Get learning rate at current step."""
        return [
            _poly_decay(
                initial_lr,
                step=step - self.hold_steps,
                decay_steps=self.max_steps - max(self.warmup_steps, self.hold_steps),
                power=self.power,
                min_lr=self.min_lr,
                cycle=self.cycle,
            )
            for initial_lr in self.base_lrs
        ]


def register_scheduler(name: str, scheduler: _LRScheduler, scheduler_params: SchedulerParams):
    """Checks if the scheduler name exists in the registry, and if it doesn't, adds it. This allows custom schedulers
    to be added and called by name during instantiation.

    Parameters
    ----------
    name : str
        Name of the optimizer. Will be used as key to retrieve the optimizer.
    scheduler : _LRScheduler
        Scheduler class (inherits from _LRScheduler)
    scheduler_params : SchedulerParams
        The parameters as a dataclass of the scheduler
    """
    if name in AVAILABLE_SCHEDULERS:
        raise ValueError(f"Cannot override pre-existing schedulers. Conflicting scheduler name = {name}")

    AVAILABLE_SCHEDULERS[name] = scheduler

    sched_name = f"{scheduler.__name__}_params"
    register_scheduler_params(name=sched_name, scheduler_params=scheduler_params)


def get_scheduler(name: str, **kwargs: Optional[Dict[str, Any]]) -> _LRScheduler:
    """Convenience method to obtain an _LRScheduler class and partially instantiate it with optimizer kwargs.

    Parameters
    ----------
    name : str
        Name of the scheduler in the registry.
    kwargs : Dict[str, Any]
        Optional kwargs of the scheduler used during instantiation.

    Returns
    -------
    _LRScheduler
        A partially instantiated _LRScheduler.
    """
    if name not in AVAILABLE_SCHEDULERS:
        raise ValueError(
            f"Cannot resolve scheduler{name}'. Available optimizers are : " f"{AVAILABLE_SCHEDULERS.keys()}"
        )

    scheduler_cls = AVAILABLE_SCHEDULERS[name]
    # Pop 'max_steps' if it's not required by the scheduler
    if "max_steps" in kwargs and "max_steps" not in inspect.signature(scheduler_cls).parameters:
        kwargs.pop("max_steps")
    scheduler = partial(scheduler_cls, **kwargs)
    return scheduler


def prepare_lr_scheduler(  # noqa: MC0001
    optimizer: optim.Optimizer,
    scheduler_config: Union[Dict[str, Any], DictConfig, None],
    train_dataloader: Optional[dataloader.DataLoader] = None,
) -> Optional[Dict[str, Any]]:
    """Constructs an LR Scheduler (optionally) for a given optimizer, based on a config with the following schema.

    Parameters
    ----------
    optimizer: The optimizer to use for the scheduler.
        name: <name of optimizer>
        lr: <maximal learning rate>
        args:
            name: auto  # special keyword, resolves to correct optimizer config for given optimizer name
            # cls: atommic.core.config.optimizers.NovogradParams  # explicit instantiation by class path
            params:  # optional override parameters for the optimizer config
                betas: [0.8, 0.5]
                weight_decay: 0.001
    scheduler_config: The scheduler config.
        name: <name of scheduler>
        iters_per_batch: null # computed at runtime; mandatory to have
        max_steps: null # computed at runtime or explicitly set here; mandatory to have
        # pytorch lightning args <mandatory>
        monitor: val_loss
        reduce_on_plateau: false
        # <scheduler config override>
        args:
            name: auto  # special keyword, resolves to correct optimizer config for given optimizer name
            # cls: atommic.core.config.schedulers.CosineAnnealingParams  # explicit instantiation by class path
            params:  # optional override parameters for the optimizer config
                warmup_steps: null
                warmup_ratio: null
                min_lr: 0.0
                last_epoch: -1

    train_dataloader: Optional requirement, must be passed if "iters_per_batch" is defined instead of "max_steps". \
    Used to compute effective "max_steps".

    Returns
    -------
    A dictionary containing the LR Scheduler implementation if the config was successfully parsed along with other \
    parameters required by Pytorch Lightning, otherwise None.
    """
    if scheduler_config is not None:
        scheduler_config = maybe_update_config_version(scheduler_config)

    # Build nested dictionary for convenience out of structured objects
    if isinstance(scheduler_config, DictConfig):
        scheduler_config = OmegaConf.to_container(scheduler_config, resolve=True)

    elif dataclasses.is_dataclass(scheduler_config):
        # Recursively transform data classes to basic dictionaries
        scheduler_config = OmegaConf.create(scheduler_config)
        scheduler_config = OmegaConf.to_container(scheduler_config, resolve=True)

    # Test to see if config follows above schema
    interval = "step"
    if scheduler_config is not None:
        if "args" in scheduler_config:
            scheduler_args = scheduler_config.pop("args")
        else:
            scheduler_args = copy.deepcopy(scheduler_config)

            # Remove extra parameters from scheduler_args nest
            # Assume all other parameters are to be passed into scheduler constructor
            scheduler_args.pop("name", None)
            scheduler_args.pop("t_max_epochs", None)
            scheduler_args.pop("t_accumulate_grad_batches", None)
            scheduler_args.pop("t_limit_train_batches", None)
            scheduler_args.pop("t_num_workers", None)
            scheduler_args.pop("monitor", None)
            scheduler_args.pop("reduce_on_plateau", None)

        if "name" in scheduler_config and scheduler_config["name"] in EPOCH_SCHEDULERS:
            interval = "epoch"
    else:
        # Return gracefully in case `sched` was not supplied; inform user
        logging.info("Scheduler not initialized as no `sched` config supplied to setup_optimizer()")
        return None

    # Try instantiation of scheduler params from config class path
    if "_target_" in scheduler_args:
        scheduler_args_cfg = OmegaConf.create(scheduler_args)
        scheduler_conf = hydra.utils.instantiate(scheduler_args_cfg)
        scheduler_args = vars(scheduler_conf)

        # Get name of the scheduler
        scheduler_name = scheduler_conf.__class__.__name__

        if "Params" in scheduler_name:
            scheduler_name = scheduler_name.replace("Params", "")

    else:
        # Class path instantiation failed; try resolving "name" component

        # Get name of the scheduler
        if "name" in scheduler_config:
            scheduler_name = scheduler_config["name"]
        else:
            logging.warning(
                "Could not resolve classpath for Scheduler Config, and `name` "
                "was not provided either. \n"
                "Scheduler cannot be instantiated !"
            )
            return None

        # If class path was not provided, perhaps `name` is provided for resolution
        if "name" in scheduler_args:
            # If `auto` is passed as name for resolution of optimizer name,
            # then lookup optimizer name and resolve its parameter config
            if scheduler_args["name"] == "auto":
                scheduler_params_name = f"{scheduler_name}Params"
            else:
                scheduler_params_name = scheduler_args["name"]

            # Get override arguments provided in the config yaml file / Dict Config
            scheduler_params_override = scheduler_args.get("params", {})

            # If params is itself a dict config object provided explicitly in Dict Config
            # Resolve to dictionary for convenience
            if isinstance(scheduler_params_override, DictConfig):
                scheduler_params_override = OmegaConf.to_container(scheduler_params_override, resolve=True)

            # Get and instantiate the Config dataclass for this scheduler
            scheduler_params_cls = get_scheduler_config(scheduler_params_name, **scheduler_params_override)
            scheduler_params = scheduler_params_cls  # instantiate the parameters object
            # extract just the dictionary from the Config object
            scheduler_args = vars(scheduler_params)

    # Extract value to monitor in losses, if provided.
    if "monitor" in scheduler_config:
        monitor = scheduler_config.get("monitor")
    else:
        # Default to train loss
        monitor = "loss"

    # Store exact max_steps if it is provided
    if "max_steps" in scheduler_config and scheduler_config["max_steps"] is not None:
        max_steps = scheduler_config["max_steps"]
    elif "t_max_epochs" in scheduler_config:
        # Compute effective max_steps if t_max_epochs is provided
        if train_dataloader is None:
            logging.warning(
                "As `t_max_epochs` is provided/computed, it is required to pass the train dataloader in order\n"
                "to compute effective maximum number of steps.\n"
                "Scheduler will not be instantiated !"
            )
            return None

        # Raise exception if neither `max_steps` nor `t_max_epochs` is provided
        if scheduler_config.get("t_max_epochs", None) is None:
            logging.warning(
                "`t_max_epochs` cannot be None when `max_steps` is not not provided.\n"
                "This can occur when `train dataloader` is not available to correctly "
                "prepare the scheduler.\n"
                "Scheduler will not be instantiated !"
            )
            return None

        # Get iters_per_batch
        max_epochs = scheduler_config.get("t_max_epochs")
        accumulate_grad_batches = scheduler_config.get("t_accumulate_grad_batches")
        limit_train_batches = scheduler_config.get("t_limit_train_batches")
        num_workers = scheduler_config.get("t_num_workers")

        # Compute effective num max_steps
        num_samples = len(train_dataloader.dataset)

        # we may need to override ModelPT setup_optimization
        if train_dataloader.batch_size is not None:
            batch_size = train_dataloader.batch_size
        elif hasattr(train_dataloader, "batch_sampler") and train_dataloader.batch_sampler is not None:
            if train_dataloader.batch_sampler.micro_batch_size is not None:
                batch_size = train_dataloader.batch_sampler.micro_batch_size
            else:
                raise ValueError(f"Could not find batch_size from batch_sampler: {train_dataloader.batch_sampler}")
        else:
            raise ValueError(f"Could not find batch_size from train_dataloader: {train_dataloader}")
        drop_last = train_dataloader.drop_last

        max_steps = compute_max_steps(
            max_epochs=max_epochs,
            accumulate_grad_batches=accumulate_grad_batches,
            limit_train_batches=limit_train_batches,
            num_workers=num_workers,
            num_samples=num_samples,
            batch_size=batch_size,
            drop_last=drop_last,
        )
    else:
        logging.warning(
            "Neither `max_steps` nor `iters_per_batch` were provided to `optim.sched`, "
            "cannot compute effective `max_steps` !\n"
            "Scheduler will not be instantiated !"
        )
        return None

    # Inject max_steps (effective or provided) into the scheduler config
    scheduler_args["max_steps"] = max_steps

    # Get the scheduler class from the config
    scheduler_cls = get_scheduler(scheduler_name, **scheduler_args)

    # Pop 'max_steps' if it's not required by the scheduler
    if "max_steps" not in inspect.signature(scheduler_cls).parameters:
        scheduler_args.pop("max_steps")

    # Instantiate the LR schedule
    schedule = scheduler_cls(optimizer, **scheduler_args)

    logging.info(
        'Scheduler "%s" \nwill be used during training (effective maximum steps = %d) - \nParameters : \n(%s)',
        str(schedule),
        max_steps,
        OmegaConf.to_yaml(OmegaConf.create(scheduler_args)),
    )

    # Wrap the schedule in PTL arguments to perform stepwise computation. Rather than epoch level computation.
    reduce_lr_on_plateau = bool(isinstance(schedule, optim.lr_scheduler.ReduceLROnPlateau))

    schedule_dict = {
        "scheduler": schedule,
        "interval": interval,
        "frequency": 1,
        "monitor": monitor,
        "reduce_on_plateau": reduce_lr_on_plateau,
    }
    return schedule_dict


def compute_max_steps(
    max_epochs, accumulate_grad_batches, limit_train_batches, num_workers, num_samples, batch_size, drop_last
):
    """Compute effective max_steps from the provided parameters.

    Parameters
    ----------
    max_epochs : int
        Maximum number of epochs to train for.
    accumulate_grad_batches : int
        Number of batches to accumulate gradients for.
    limit_train_batches : int
        Number of batches to train for.
    num_workers : int
        Number of workers to use for training.
    num_samples : int
        Number of samples in the dataset.
    batch_size : int
        Batch size.
    drop_last : bool
        Whether to drop the last batch or not.
    """
    _round = math.floor if drop_last else math.ceil

    sampler_num_samples = math.ceil(num_samples / max(1, num_workers))

    if drop_last and num_workers > 1:
        logging.warning(
            "Please note that drop_last is broken in pytorch 1.6.0. We will fix when pytorch 1.7.0 is released"
        )
        # TODO: Master version, not in pytorch 1.6.0

    steps_per_epoch = _round(sampler_num_samples / batch_size)
    if isinstance(limit_train_batches, int) or limit_train_batches == 0.0:
        steps_per_epoch = min(steps_per_epoch, int(limit_train_batches))
    elif steps_per_epoch != float("inf"):
        # limit_train_batches is a percentage of batches per epoch
        steps_per_epoch = int(steps_per_epoch * limit_train_batches)

    return math.ceil(steps_per_epoch / accumulate_grad_batches) * max_epochs


AVAILABLE_SCHEDULERS = {
    "WarmupPolicy": WarmupPolicy,
    "WarmupHoldPolicy": WarmupHoldPolicy,
    "SquareAnnealing": SquareAnnealing,
    "CosineAnnealing": CosineAnnealing,
    "NoamAnnealing": NoamAnnealing,
    "NoamHoldAnnealing": NoamHoldAnnealing,
    "WarmupAnnealing": WarmupAnnealing,
    "InverseSquareRootAnnealing": InverseSquareRootAnnealing,
    "T5InverseSquareRootAnnealing": T5InverseSquareRootAnnealing,
    "SquareRootAnnealing": SquareRootAnnealing,
    "PolynomialDecayAnnealing": PolynomialDecayAnnealing,
    "PolynomialHoldDecayAnnealing": PolynomialHoldDecayAnnealing,
    "StepLR": pt_scheduler.StepLR,
    "ExponentialLR": pt_scheduler.ExponentialLR,
    "ReduceLROnPlateau": pt_scheduler.ReduceLROnPlateau,
    "CyclicLR": pt_scheduler.CyclicLR,
}

EPOCH_SCHEDULERS = {
    "ExponentialLR": pt_scheduler.ExponentialLR,
    "ReduceLROnPlateau": pt_scheduler.ReduceLROnPlateau,
}
