"""Some helper functions for machine learning tasks."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_utils/40_ml_utils.ipynb.

# %% auto 0
__all__ = ['LRSchedulerPerStep']

# %% ../nbs/00_utils/40_ml_utils.ipynb 3
from typing import  List, Tuple, Literal
import torch

# %% ../nbs/00_utils/40_ml_utils.ipynb 4
class LRSchedulerPerStep():
    """
    Learning rate scheduler from Attention is all you need paper (https://arxiv.org/abs/1706.03762)
    One ajustment: Added base LR as tunable parameter rather than setting it automated based on model dimension
    """
    
    def __init__(self,
                optimizer: torch.optim.Optimizer, # Optimizer to adjust learning rate for
                base_learning_rate: float = 0.0001,
                warmup: int =4000):

        # Ensure optimizer is a PyTorch optimizer
        if not isinstance(optimizer, torch.optim.Optimizer):
            raise ValueError('Optimizer must be a PyTorch optimizer')
        
        self.optimizer = optimizer
        self.basic = base_learning_rate
        self.warm = warmup**-1.5
        self.scaling_factor = 1/warmup**-0.5 # ensures that the peak realtive to the base lr is always 1

        self.step_num = 0   
        self.step()
        
    def step(self):
        self.step_num += 1
        lr = self.basic * self.scaling_factor * min(self.step_num**-0.5, self.step_num*self.warm)
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
