import torch
from torch.optim.lr_scheduler import _LRScheduler
import math


class CosineLR(_LRScheduler):
    def __init__(self, optimizer, init_lr, total_epochs, last_epoch=-1):
        super(CosineLR, self).__init__(optimizer, last_epoch=-1)
        self.optimizer = optimizer
        self.init_lr = init_lr
        self.total_epochs = total_epochs
        self.last_epoch = last_epoch
        print(f'CosineLR start from epoch(step) {last_epoch} with init_lr {init_lr} ')

    def get_lr(self):
        if self.last_epoch == 0:
            return [group['lr'] for group in self.optimizer.param_groups]

        return [0.5 * (1 + math.cos(self.last_epoch * math.pi / self.total_epochs)) * self.init_lr for group in
                self.optimizer.param_groups]