"""Custom WSD learning rate schedulers for composer."""


from composer.core import State, Time
from composer.optim.scheduler import (
    ComposerScheduler,
    LinearScheduler,
    SqrtScheduler,
    _convert_time,  # noqa: PLC2701
    _raise_if_cooldown_and_max_incompatible,  # noqa: PLC2701
    _raise_if_cooldown_plus_warmup_and_max_incompatible,  # noqa: PLC2701
    _raise_if_max_duration_exceeds_t_max,  # noqa: PLC2701
    _raise_if_warmup_and_max_incompatible,  # noqa: PLC2701
)
from llmfoundry.registry import schedulers

__all__ = [
    "ConstantWithLinearCooldownWithWarmupScheduler",
    "ConstantWithSqrtCooldownWithWarmupScheduler",
]


@schedulers.register_class("constant_with_linear_cooldown_with_warmup")
class ConstantWithLinearCooldownWithWarmupScheduler(ComposerScheduler):
    """WSD Scheduler with Linear Cooldown.

    A scheduler that maintains a constant learning rate with optional linear warmup
    and cooldown periods.

    Inspired by https://arxiv.org/abs/2405.18392v3.

    Atributes
    ---------
        t_warmup :  Union[str, Time]
            The duration of the warmup period.
        t_cooldown :  Union[str, Time]
            The duration of the cooldown period.
        t_max Union[str, Time]:
            The total duration of the warmup and cooldown periods.
        scale_warmup : bool:
            If True, scales the learning rate during the warmup period.
        scale_cooldown : bool
            If True, scales the learning rate during the cooldown period.
        warmup_scheduler : LinearScheduler
            The scheduler used during the warmup period.
        cooldown_scheduler :  LinearScheduler
            The scheduler used during the cooldown period.
    """

    def __init__(
        self,
        t_warmup: str | Time,
        t_cooldown: str | Time,
        t_max: str | Time = "1dur",
        *,
        scale_warmup: bool = False,
        scale_cooldown: bool = False,
    ) -> None:
        """Initialise WSD scheduler with linear cooldown.

        Parameters
        ----------
        t_warmup : Union[str, Time]
            The duration of the warmup period.
        t_cooldown : Union[str, Time]
            The duration of the cooldown period.
        t_max : Union[str, Time], optional
            The total duration of the warmup and cooldown periods, by default "1dur"
        scale_warmup : bool, optional
            If True, scale learning rate during the warmup period, by default False
        scale_cooldown : bool, optional
            If True, scale learning rate during the cooldown period, by default False

        """
        self.t_warmup = t_warmup
        self.t_cooldown = t_cooldown
        self.t_max = t_max
        self.scale_warmup = scale_warmup
        self.scale_cooldown = scale_cooldown
        self.warmup_scheduler = LinearScheduler(
            alpha_i=0.0,
            alpha_f=1.0,
            t_max=t_warmup,
        )
        self.cooldown_scheduler = LinearScheduler(
            alpha_i=1.0,
            alpha_f=0.0,
            t_max=t_cooldown,
        )

    def __call__(self, state: State, ssr: float = 1.0) -> float:
        """Calculate learning rate multiplier based on the current state and schedule.

        Parameters
        ----------
        state : State
            The current state of training.
        ssr : float, optional
            The scale factor for the learning rate, by default 1.0

        Returns
        -------
        float
           The learning rate multiplier.

        """
        assert state.max_duration is not None, (
            "max_duration should be set whenever schedulers are invoked"
        )

        # Convert warmup, cooldown, and max durations to the appropriate time units
        t_warmup = _convert_time(self.t_warmup, state)
        t_cooldown = _convert_time(self.t_cooldown, state)
        t_max = _convert_time(self.t_max, state, ssr=ssr)

        # Raise errors if warmup, cooldown, or max durations are incompatible
        _raise_if_warmup_and_max_incompatible(t_warmup, t_max)
        _raise_if_cooldown_and_max_incompatible(t_cooldown, t_max)
        _raise_if_cooldown_plus_warmup_and_max_incompatible(
            t_cooldown + t_warmup,
            t_max,
        )
        _raise_if_max_duration_exceeds_t_max(t_max, state)

        # If within the warmup period, use the warmup scheduler
        if state.timestamp < t_warmup:
            if self.scale_warmup:
                return self.warmup_scheduler(state, ssr)
            return self.warmup_scheduler(state)

        # If within the cooldown period, use the cooldown scheduler
        if state.timestamp >= t_max - t_cooldown:
            if self.scale_cooldown:
                return self.cooldown_scheduler(state, ssr)
            return self.cooldown_scheduler(state)

        # Otherwise, return a constant learning rate multiplier of 1.0
        return 1.0


@schedulers.register_class("constant_with_sqrt_cooldown_with_warmup")
class ConstantWithSqrtCooldownWithWarmupScheduler(ComposerScheduler):
    """WSD Scheduler with (1-sqrt) cooldown.

    A scheduler that maintains a constant learning rate with optional linear warmup
    and (1-sqrt) cooldown periods.

    Inspired by https://arxiv.org/abs/2405.18392v3.

    Atributes
    ---------
        t_warmup :  Union[str, Time]
            The duration of the warmup period.
        t_cooldown :  Union[str, Time]
            The duration of the cooldown period.
        t_max Union[str, Time]:
            The total duration of the warmup and cooldown periods.
        scale_warmup : bool:
            If True, scales the learning rate during the warmup period.
        scale_cooldown : bool
            If True, scales the learning rate during the cooldown period.
        warmup_scheduler : LinearScheduler
            The scheduler used during the warmup period.
        cooldown_scheduler :  LinearScheduler
            The scheduler used during the cooldown period.
    """

    def __init__(  # noqa: PLR0913
        self,
        t_warmup: str | Time,
        t_cooldown: str | Time,
        t_max: str | Time = "1dur",
        alpha_f: float = 1.0,
        *,
        scale_warmup: bool = False,
        scale_cooldown: bool = False,
    ) -> None:
        """Initialise WSD scheduler with (1-sqrt) cooldown.

        Parameters
        ----------
        t_warmup : Union[str, Time]
            The duration of the warmup period.
        t_cooldown : Union[str, Time]
            The duration of the cooldown period.
        t_max : Union[str, Time], optional
            The total duration of the warmup and cooldown periods, by default "1dur"
        alpha_f : bool
            Final learning rate multiplier, by default 1.0
        scale_warmup : bool, optional
            If True, scale learning rate during the warmup period, by default False
        scale_cooldown : bool, optional
            If True, scale learning rate during the cooldown period, by default False

        """
        self.t_warmup = t_warmup
        self.t_cooldown = t_cooldown
        self.t_max = t_max
        self.scale_warmup = scale_warmup
        self.scale_cooldown = scale_cooldown
        self.warmup_scheduler = LinearScheduler(
            alpha_i=0.0,
            alpha_f=alpha_f,
            t_max=t_warmup,
        )
        self.cooldown_scheduler = SqrtScheduler(t_max=t_max, t_duration=t_cooldown)

    def __call__(self, state: State, ssr: float = 1.0) -> float:
        """Calculate learning rate multiplier based on  current state and schedule.

        Parameters
        ----------
        state : State
            The current state of training.
        ssr : float, optional
            The scale factor for the learning rate, by default 1.0

        Returns
        -------
        float
            _desThe learning rate multiplier.cription_

        """
        assert state.max_duration is not None, (
            "max_duration should be set whenever schedulers are invoked"
        )

        # Convert warmup, cooldown, and max durations to the appropriate time units
        t_warmup = _convert_time(self.t_warmup, state)
        t_cooldown = _convert_time(self.t_cooldown, state)
        t_max = _convert_time(self.t_max, state, ssr=ssr)

        # Raise errors if warmup, cooldown, or max durations are incompatible
        _raise_if_warmup_and_max_incompatible(t_warmup, t_max)
        _raise_if_cooldown_and_max_incompatible(t_cooldown, t_max)
        _raise_if_cooldown_plus_warmup_and_max_incompatible(
            t_cooldown + t_warmup,
            t_max,
        )
        _raise_if_max_duration_exceeds_t_max(t_max, state)

        # If within the warmup period, use the warmup scheduler
        if state.timestamp < t_warmup:
            if self.scale_warmup:
                return self.warmup_scheduler(state, ssr)
            return self.warmup_scheduler(state)

        # If within the cooldown period, use the cooldown scheduler
        if state.timestamp >= t_max - t_cooldown:
            if self.scale_cooldown:
                return self.cooldown_scheduler(state, ssr)
            return self.cooldown_scheduler(state)

        # Otherwise, return a constant learning rate multiplier of 1.0
        return 1.0


@schedulers.register_class("constant_with_sqrt_cooldown_with_warmup_and_switch")
class ConstantWithSqrtCooldownWithWarmupAndSwitchScheduler(ComposerScheduler):
    """WSD Scheduler with (1-sqrt) cooldown.

    A scheduler that maintains a constant learning rate with optional linear warmup
    and (1-sqrt) cooldown periods.

    Inspired by https://arxiv.org/abs/2405.18392v3.

    Atributes
    ---------
        t_warmup :  Union[str, Time]
            The duration of the warmup period.
        t_cooldown :  Union[str, Time]
            The duration of the cooldown period.
        t_max Union[str, Time]:
            The total duration of the warmup and cooldown periods.
        scale_warmup : bool:
            If True, scales the learning rate during the warmup period.
        scale_cooldown : bool
            If True, scales the learning rate during the cooldown period.
        warmup_scheduler : LinearScheduler
            The scheduler used during the warmup period.
        cooldown_scheduler :  LinearScheduler
            The scheduler used during the cooldown period.
    """

    def __init__(  # noqa: PLR0913
        self,
        t_warmup: str | Time,
        t_cooldown: str | Time,
        switch_scale: float = 1.0,
        t_max: str | Time = "1dur",
        *,
        scale_warmup: bool = False,
        scale_cooldown: bool = False,
    ) -> None:
        """Initialise WSD scheduler with (1-sqrt) cooldown.

        Parameters
        ----------
        t_warmup : Union[str, Time]
            The duration of the warmup period.
        t_cooldown : Union[str, Time]
            The duration of the cooldown period.
        switch_scale: float
            The scale factor for switching, by default 1.0
        t_max : Union[str, Time], optional
            The total duration of the warmup and cooldown periods, by default "1dur"
        scale_warmup : bool, optional
            If True, scale learning rate during the warmup period, by default False
        scale_cooldown : bool, optional
            If True, scale learning rate during the cooldown period, by default False

        """
        self.t_warmup = t_warmup
        self.t_cooldown = t_cooldown
        self.t_max = t_max
        self.scale_warmup = scale_warmup
        self.scale_cooldown = scale_cooldown
        self.warmup_scheduler = LinearScheduler(
            alpha_i=0.0,
            alpha_f=1.0,
            t_max=t_warmup,
        )
        self.cooldown_scheduler = SqrtScheduler(t_max=t_max, t_duration=t_cooldown)
        self.switch_scale = switch_scale

    def __call__(self, state: State, ssr: float = 1.0) -> float:
        """Calculate learning rate multiplier based on  current state and schedule.

        Parameters
        ----------
        state : State
            The current state of training.
        ssr : float, optional
            The scale factor for the learning rate, by default 1.0

        Returns
        -------
        float
            _desThe learning rate multiplier.cription_

        """
        assert state.max_duration is not None, (
            "max_duration should be set whenever schedulers are invoked"
        )

        # Convert warmup, cooldown, and max durations to the appropriate time units
        t_warmup = _convert_time(self.t_warmup, state)
        t_cooldown = _convert_time(self.t_cooldown, state)
        t_max = _convert_time(self.t_max, state, ssr=ssr)

        # Raise errors if warmup, cooldown, or max durations are incompatible
        _raise_if_warmup_and_max_incompatible(t_warmup, t_max)
        _raise_if_cooldown_and_max_incompatible(t_cooldown, t_max)
        _raise_if_cooldown_plus_warmup_and_max_incompatible(
            t_cooldown + t_warmup,
            t_max,
        )
        _raise_if_max_duration_exceeds_t_max(t_max, state)

        # If within the warmup period, use the warmup scheduler
        if state.timestamp < t_warmup:
            if self.scale_warmup:
                return self.warmup_scheduler(state, ssr)
            return self.warmup_scheduler(state)

        # If within the cooldown period, use the cooldown scheduler
        if state.timestamp >= t_max - t_cooldown:
            if self.scale_cooldown:
                return self.cooldown_scheduler(state, ssr)
            return self.cooldown_scheduler(state) * self.switch_scale

        # Otherwise, return a constant learning rate multiplier of 1.0
        return self.switch_scale
