import math
from torch.optim.lr_scheduler import LambdaLR

def cosine_lr(optimizer, base_lr, warmup_length, total_steps):
    import math

    # ✅ 如果 warmup_length 是 list，则取第一个值
    if isinstance(warmup_length, (list, tuple)):
        warmup_length = warmup_length[0]
    if isinstance(total_steps, (list, tuple)):
        total_steps = total_steps[0]

    def lr_lambda(current_step):
        if current_step < warmup_length:
            return float(current_step) / float(max(1, warmup_length))
        progress = float(current_step - warmup_length) / float(max(1, total_steps - warmup_length))
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    return LambdaLR(optimizer, lr_lambda)