from functools import partial
from fvcore.common.param_scheduler import MultiStepParamScheduler

from detectron2 import model_zoo
from detectron2.config import LazyCall as L
from detectron2.config import CfgNode, LazyConfig
from detectron2.solver import WarmupParamScheduler

from detectron2.modeling.backbone.vit import get_vit_lr_decay_rate
import os
from ..common.coco_loader_lsj import dataloader

# model = model_zoo.get_config("common/models/mask_rcnn_vitdet.py").model
# model.backbone.square_pad = 512 # change input size to 512x512

cfg_file = "./models/mask_rcnn_cwm.py"
model = LazyConfig.load(cfg_file).model

# Initialization and trainer settings
train = model_zoo.get_config("common/train.py").train
train.amp.enabled = True
train.ddp.fp16_compression = True
train.init_checkpoint = (
    '/home/honglinc/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth'
)
train.output_dir = '/ccn2/u/honglinc/cwm_checkpoints/dinov2_coco_finetune_512'

# model.backbone.net.window_size = 0
# model.backbone.net.window_block_indexes = []
# model.backbone.net.use_rel_pos = False
# model.backbone.net.drop_path_rate = 0.

# Schedule
# 100 ep = 184375 iters * 64 images/iter / 118000 images/ep
# 100 ep = 184375 iters * 64 images/iter / 118000 images/ep
# train.max_iter = 184375
# milestones = [163889, 177546]

# 50 ep = 30730 iters * 96 images/iter / 118000 images/ep
train.max_iter = 61458
milestones = [54629, 59182]

lr_multiplier = L(WarmupParamScheduler)(
    scheduler=L(MultiStepParamScheduler)(
        values=[1.0, 0.1, 0.01],
        milestones=milestones,
        num_updates=train.max_iter,
    ),
    warmup_length=250 / train.max_iter,
    warmup_factor=0.001,
)

# Optimizer
optimizer = model_zoo.get_config("common/optim.py").AdamW
optimizer.params.lr_factor_func = partial(get_vit_lr_decay_rate, num_layers=12, lr_decay_rate=0.7)
optimizer.params.overrides = {"pos_embed": {"weight_decay": 0.0}}