from ganime.trainer.warmup.cosine import WarmUpCosine


def create_warmup_scheduler(trainer_config, num_devices):
    len_x_train = trainer_config["len_x_train"]
    batch_size = trainer_config["batch_size"]
    n_epochs = trainer_config["n_epochs"]

    total_steps = int(len_x_train / batch_size * n_epochs / num_devices)
    warmup_epoch_percentage = trainer_config["warmup_epoch_percentage"]
    warmup_steps = int(total_steps * warmup_epoch_percentage)

    scheduled_lrs = WarmUpCosine(
        lr_start=trainer_config["lr_start"],
        lr_max=trainer_config["lr_max"],
        warmup_steps=warmup_steps,
        total_steps=total_steps,
    )

    return scheduled_lrs