import math

class LinearWarmupCosineLRScheduler:
    def __init__(
        self,
        optimizer,
        min_lr_list,
        init_lr_list,
        warmup_steps=0,
        warmup_start_lr_list=None,
        **kwargs
    ):
        self.optimizer = optimizer

        self.min_lr_list = min_lr_list
        self.init_lr_list = init_lr_list
        self.warmup_steps = warmup_steps
        self.warmup_start_lr_list = warmup_start_lr_list if warmup_start_lr_list is not None else init_lr_list

    def step(self, cur_step, cur_epoch, max_step):
        for i, param_group in enumerate(self.optimizer.param_groups):
            if cur_epoch == 0 and cur_step < self.warmup_steps:
                lr = self.warmup_lr_schedule(cur_step, self.warmup_start_lr_list[i], self.init_lr_list[i])
            else:
                lr = self.cosine_lr_schedule(cur_step - self.warmup_steps, max_step - self.warmup_steps, self.init_lr_list[i], self.min_lr_list[i])
            param_group["lr"] = lr
            
    def cosine_lr_schedule(self, step, max_step, init_lr, min_lr):
        """Decay the learning rate using cosine schedule"""
        lr = (init_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * step / max_step)) + min_lr
        return lr

    def warmup_lr_schedule(self, step, init_lr, max_lr):
        """Warmup the learning rate"""
        lr = min(max_lr, init_lr + (max_lr - init_lr) * step / max(self.warmup_steps, 1))
        return lr

    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)