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

class CosineAnnealingWithWarmupScheduler(_LRScheduler):
    def __init__(self, optimizer, t_warmup, t_max, alpha_f=0.0, scale_warmup=False, last_epoch=-1):
        self.t_warmup = t_warmup
        self.t_max = t_max
        self.alpha_f = alpha_f
        self.scale_warmup = scale_warmup
        super(CosineAnnealingWithWarmupScheduler, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch < self.t_warmup:
            # Warmup phase
            return [base_lr * (self.last_epoch / self.t_warmup) for base_lr in self.base_lrs]
        else:
            # Cosine annealing phase
            cos_inner = (1 + math.cos(math.pi * (self.last_epoch - self.t_warmup) / (self.t_max - self.t_warmup))) / 2
            return [base_lr * (self.alpha_f + (1 - self.alpha_f) * cos_inner) for base_lr in self.base_lrs]

class WarmupScheduler(_LRScheduler):
    def __init__(self, optimizer, t_warmup, last_epoch=-1):
        self.t_warmup = t_warmup
        super(WarmupScheduler, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch < self.t_warmup:
            # Warmup phase
            return [base_lr * (self.last_epoch / self.t_warmup) for base_lr in self.base_lrs]
        else:
            return [base_lr for base_lr in self.base_lrs]