from detectron2.config import LazyCall as L
from detectron2.solver import WarmupParamScheduler
from fvcore.common.param_scheduler import MultiStepParamScheduler


def default_X_scheduler(num_X):
    """
    Returns the config for a default multi-step LR scheduler such as "1x", "3x",
    commonly referred to in papers, where every 1x has the total length of 1440k
    training images (~12 COCO epochs). LR is decayed twice at the end of training
    following the strategy defined in "Rethinking ImageNet Pretraining", Sec 4.
    Args:
        num_X: a positive real number
    Returns:
        DictConfig: configs that define the multiplier for LR during training
    """
    # total number of iterations assuming 16 batch size, using 1440000/16=90000
    total_steps_16bs = num_X * 90000

    if num_X <= 2:
        scheduler = L(MultiStepParamScheduler)(
            values=[1.0, 0.1, 0.01],
            # note that scheduler is scale-invariant. This is equivalent to
            # milestones=[6, 8, 9]
            milestones=[60000, 80000, 90000],
        )
    else:
        scheduler = L(MultiStepParamScheduler)(
            values=[1.0, 0.1, 0.01],
            milestones=[
                total_steps_16bs - 60000,
                total_steps_16bs - 20000,
                total_steps_16bs,
            ],
        )
    return L(WarmupParamScheduler)(
        scheduler=scheduler,
        warmup_length=1000 / total_steps_16bs,
        warmup_method="linear",
        warmup_factor=0.001,
    )


def default_coco_scheduler(epochs=50, decay_epochs=40, warmup_epochs=0):
    """
    Returns the config for a default multi-step LR scheduler such as "50epochs",
    commonly referred to in papers, where every 1x has the total length of 1440k
    training images (~12 COCO epochs). LR is decayed once at the end of training.

    Args:
        epochs (int): total training epochs.
        decay_epochs (int): lr decay steps.
        warmup_epochs (int): warmup epochs.

    Returns:
        DictConfig: configs that define the multiplier for LR during training
    """
    # total number of iterations assuming 16 batch size, using 1440000/16=90000
    total_steps_16bs = epochs * 7500
    decay_steps = decay_epochs * 7500
    warmup_steps = warmup_epochs * 7500
    scheduler = L(MultiStepParamScheduler)(
        values=[1.0, 0.1],
        milestones=[decay_steps, total_steps_16bs],
    )
    return L(WarmupParamScheduler)(
        scheduler=scheduler,
        warmup_length=warmup_steps / total_steps_16bs,
        warmup_method="linear",
        warmup_factor=0.001,
    )


def multi_steps_scheduler(
    epochs=12,
    decay_epochs=[8, 11],
    warmup_steps=1000,
    samples_per_epoch=100000,
    batch_size=16,
    _total_steps=-1,
):
    """
    Returns the config for a default multi-step LR scheduler such as "50epochs",
    commonly referred to in papers, where every 1x has the total length of 1440k
    training images (~12 COCO epochs). LR is decayed once at the end of training.

    Args:
        epochs (int): total training epochs.
        decay_epochs (int): lr decay steps.
        warmup_steps (int): warmup steps.

    Returns:
        DictConfig: configs that define the multiplier for LR during training
    """
    # total number of iterations assuming 16 batch size, using 1440000/16=90000
    steps_per_epoch = int(samples_per_epoch / batch_size)
    total_steps = epochs * steps_per_epoch
    total_steps = max(total_steps, _total_steps)
    decay_steps = [decay_epoch * steps_per_epoch for decay_epoch in decay_epochs]
    decay_steps.append(total_steps)
    values = [0.1**i for i in range(len(decay_epochs) + 1)]
    scheduler = L(MultiStepParamScheduler)(
        values=values,
        milestones=decay_steps,
    )
    return L(WarmupParamScheduler)(
        scheduler=scheduler,
        warmup_length=warmup_steps / total_steps,
        warmup_method="linear",
        warmup_factor=0.001,
    )


# default coco scheduler
lr_multiplier_1x = default_X_scheduler(1)
lr_multiplier_2x = default_X_scheduler(2)
lr_multiplier_3x = default_X_scheduler(3)
lr_multiplier_6x = default_X_scheduler(6)
lr_multiplier_9x = default_X_scheduler(9)


# default scheduler for detr
lr_multiplier_50ep = default_coco_scheduler(50, 40, 0)
lr_multiplier_36ep = default_coco_scheduler(36, 30, 0)
lr_multiplier_24ep = default_coco_scheduler(24, 20, 0)
lr_multiplier_12ep = default_coco_scheduler(12, 11, 0)

# warmup scheduler for detr
lr_multiplier_50ep_warmup = default_coco_scheduler(50, 40, 1e-3)
lr_multiplier_12ep_warmup = default_coco_scheduler(12, 11, 1e-3)

# warmup scheduler for ovdino on COCO fine-tuning
# lr_multiplier_110k_bs16_16ep_two_steps_warmup = multi_steps_scheduler(
#     16, [12, 14], 1000, 114000, 16, 114000
# )
lr_multiplier_75k_bs16_16ep_two_steps_warmup = multi_steps_scheduler(
    16, [8, 12], 1000, 75000, 16, 75000
)

lr_multiplier_14k_bs16_16ep_two_steps_warmup_ArTaxOr = multi_steps_scheduler(
    16, [12, 14], 1000, 14000, 16, 14000
)



lr_multiplier_14k_bs32_24ep_two_steps_warmup_ArTaxOr = multi_steps_scheduler(
    24, [16, 20], 500, 14000, 32, 10500
)

lr_multiplier_14k_bs32_16ep_two_steps_warmup_ArTaxOr = multi_steps_scheduler(
    16, [12, 14], 1000, 14000, 32, 7000
)

lr_multiplier_0p5k_bs16_16ep_two_steps_warmup_clipart1k = multi_steps_scheduler(
    16, [12, 14], 100, 500, 16, 500
)
lr_multiplier_120k_bs32_24ep_two_steps_warmup = multi_steps_scheduler(
    24, [16, 20], 1000, 120000, 32, 90000
)

# def multi_steps_scheduler(
#     epochs=12,
#     decay_epochs=[8, 11],
#     warmup_steps=1000,
#     samples_per_epoch=100000,
#     batch_size=16,
#     _total_steps=-1,
# ):

# lr_multiplier_14k_bs16_24ep_two_steps_warmup_clipart1k = multi_steps_scheduler(
#     24, [16, 20], 100, 50, 16, 750
# )

lr_multiplier_14k_bs16_24ep_two_steps_warmup_ArTaxOr = multi_steps_scheduler(
    24, [16, 20], 1000, 14000, 16, 21000
)

lr_multiplier_14k_bs16_24ep_two_steps_warmup_dior = multi_steps_scheduler(
    24, [16, 20], 1000, 18500, 16, 27750
)

lr_multiplier_14k_bs16_24ep_two_steps_warmup_neudet = multi_steps_scheduler(
    24, [16, 20], 1000, 1500, 16, 2550
)

lr_multiplier_14k_bs16_24ep_two_steps_warmup_uodd = multi_steps_scheduler(
    24, [16, 20], 1000, 2700, 16, 4050
)