"""PyTorch optimizer schdulers."""

from typing import List

from torch.optim import Optimizer

from espnet.scheduler.scheduler import SchedulerInterface


class PyTorchScheduler:
    """PyTorch optimizer scheduler."""

    def __init__(self, schedulers: List[SchedulerInterface], optimizer: Optimizer):
        """Initialize class."""
        self.schedulers = schedulers
        self.optimizer = optimizer
        for s in self.schedulers:
            for group in optimizer.param_groups:
                group.setdefault("initial_" + s.key, group[s.key])

    def step(self, n_iter: int):
        """Update optimizer by scheduling."""
        for s in self.schedulers:
            for group in self.optimizer.param_groups:
                group[s.key] = group["initial_" + s.key] * s.scale(n_iter)