from __future__ import annotations

from typing import Any, Callable, Optional

import numpy as np
import numpy.typing as npt
from typing_extensions import Literal, TypedDict

from .._typing import Tensor

class LRStateDict(TypedDict):
    last_epoch: int
    last_lr: float

class LinearWarmupLRStateDict(LRStateDict):
    LinearWarmup_LR: LRStateDict

class ReduceOnPlateauLRStateDict(LRStateDict):
    cooldown_counter: int
    best: int
    num_bad_epochs: int

class LRScheduler:
    base_lr: Any = ...
    last_lr: Any = ...
    last_epoch: Any = ...
    verbose: Any = ...
    def __init__(
        self,
        learning_rate: float = ...,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def __call__(self) -> float: ...
    def step(self, epoch: Optional[int] = ...) -> None: ...
    def state_dict(self) -> LRStateDict: ...
    keys: str = ...
    def state_keys(self) -> None: ...
    def set_state_dict(self, state_dict: LRStateDict) -> None: ...
    set_dict = set_state_dict
    def get_lr(self) -> float: ...

class NoamDecay(LRScheduler):
    d_model: Any = ...
    warmup_steps: Any = ...
    def __init__(
        self,
        d_model: int,
        warmup_steps: int,
        learning_rate: float = ...,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class PiecewiseDecay(LRScheduler):
    boundaries: Any = ...
    values: Any = ...
    def __init__(
        self,
        boundaries: list[int] | tuple[int, ...],
        values: list[float] | tuple[float, ...],
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class NaturalExpDecay(LRScheduler):
    gamma: Any = ...
    def __init__(
        self,
        learning_rate: float,
        gamma: float,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class InverseTimeDecay(LRScheduler):
    gamma: Any = ...
    def __init__(
        self,
        learning_rate: float,
        gamma: float,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class PolynomialDecay(LRScheduler):
    decay_steps: Any = ...
    end_lr: Any = ...
    power: Any = ...
    cycle: Any = ...
    def __init__(
        self,
        learning_rate: float,
        decay_steps: int,
        end_lr: float = ...,
        power: float = ...,
        cycle: bool = ...,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class LinearWarmup(LRScheduler):
    learning_rate: Any = ...
    warmup_steps: Any = ...
    start_lr: Any = ...
    end_lr: Any = ...
    def __init__(
        self,
        learning_rate: float | LRScheduler,
        warmup_steps: int,
        start_lr: float,
        end_lr: float,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def state_dict(self) -> LinearWarmupLRStateDict | LRStateDict: ...
    def set_state_dict(
        self,
        state_dict: LinearWarmupLRStateDict | LRStateDict,
    ) -> None: ...
    def get_lr(self) -> float: ...

class ExponentialDecay(LRScheduler):
    gamma: Any = ...
    def __init__(
        self,
        learning_rate: float,
        gamma: float,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class MultiStepDecay(LRScheduler):
    milestones: Any = ...
    gamma: Any = ...
    def __init__(
        self,
        learning_rate: float,
        milestones: list[int] | tuple[int, ...],
        gamma: float = ...,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class StepDecay(LRScheduler):
    step_size: Any = ...
    gamma: Any = ...
    def __init__(
        self,
        learning_rate: float,
        step_size: int,
        gamma: float = ...,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class LambdaDecay(LRScheduler):
    lr_lambda: Any = ...
    def __init__(
        self,
        learning_rate: float,
        lr_lambda: Callable[[float], float],
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class ReduceOnPlateau(LRScheduler):
    mode: Any = ...
    factor: Any = ...
    threshold_mode: Any = ...
    patience: Any = ...
    threshold: Any = ...
    cooldown: Any = ...
    min_lr: Any = ...
    epsilon: Any = ...
    cooldown_counter: int = ...
    best: Any = ...
    num_bad_epochs: int = ...
    base_lr: Any = ...
    last_lr: Any = ...
    last_epoch: int = ...
    verbose: Any = ...
    def __init__(
        self,
        learning_rate: float,
        mode: Literal["min", "max"] = ...,
        factor: float = ...,
        patience: int = ...,
        threshold: float = ...,
        threshold_mode: Literal["rel", "abs"] = ...,
        cooldown: int = ...,
        min_lr: float = ...,
        epsilon: float = ...,
        verbose: bool = ...,
    ) -> None: ...
    keys: Any = ...
    def state_keys(self) -> None: ...
    def step(
        self,
        metrics: Tensor | npt.NDArray[Any] | float,
        epoch: Optional[int] = ...,
    ) -> None: ...

class CosineAnnealingDecay(LRScheduler):
    T_max: Any = ...
    eta_min: Any = ...
    def __init__(
        self,
        learning_rate: float,
        T_max: int,
        eta_min: float | int = ...,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...

class MultiplicativeDecay(LRScheduler):
    lr_lambda: Any = ...
    def __init__(
        self,
        learning_rate: float,
        lr_lambda: Callable[[float], float],
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self) -> float: ...
