import copy
from typing import List

import torch

class NoamScheduler:
    """OG scheduler from transformer paper: https://arxiv.org/pdf/1706.03762.pdf
    Implementation from Annotated Transformer: https://nlp.seas.harvard.edu/2018/04/03/attention.html
    """

    def __init__(
        self,
        optimizer: torch.optim.Optimizer,
        d_model: int = 512,
        factor: float = 1.0,
        warmup: int = 4000,
    ):
        # Store hparams
        self.warmup = warmup
        self.factor = factor
        self.d_model = d_model

        # Initialize variables `lr` and `steps`
        self.lr = None
        self.steps = 0

        # Store the optimizer
        self.optimizer = optimizer

    def state_dict(self):
        return {
            key: value for key, value in self.__dict__.items() if key != "optimizer"
        }

    def load_state_dict(self, state_dict):
        self.__dict__.update(state_dict)

    def step(self):
        self.steps += 1
        self.lr = self.factor * (
            self.d_model ** (-0.5)
            * min(self.steps ** (-0.5), self.steps * self.warmup ** (-1.5))
        )

        for p in self.optimizer.param_groups:
            p["lr"] = self.lr