from detrex.config import get_config
from ..models.dino_r50 import model

import itertools

from omegaconf import OmegaConf

from detectron2.config import LazyCall as L
from detectron2.data import (
    build_detection_test_loader,
    build_detection_train_loader,
    get_detection_dataset_dicts,
)
from detectron2.data.datasets import register_coco_instances

from detectron2.modeling.backbone import ResNet, BasicStem

from projects.vCLR_deformable_mask.modeling import OursDatasetMapper

dataloader = OmegaConf.create()




# get default config
optimizer = get_config("common/optim.py").AdamW
lr_multiplier = get_config("common/coco_schedule.py").lr_multiplier_12ep
train = get_config("common/train.py").train

# modify training config
train.init_checkpoint = "detectron2://ImageNetPretrained/torchvision/R-50.pkl"
train.output_dir = "./output/dino_openworld"

# max training iterations
train.max_iter = 60000
train.eval_period = 5000
train.log_period = 200
train.checkpointer.period = 5000

# gradient clipping for training
train.clip_grad.enabled = True
train.clip_grad.params.max_norm = 0.1
train.clip_grad.params.norm_type = 2

# set training devices
train.device = "cuda"
model.device = train.device

# modify optimizer config
optimizer.lr = 1e-4 # original 1e-4
optimizer.betas = (0.9, 0.999)
optimizer.weight_decay = 1e-4
optimizer.params.lr_factor_func = lambda module_name: 0.1 if "backbone" in module_name else 1




# modify model config
model.dn_number = 100
model.num_classes = 1
model.select_box_nums_for_evaluation=900

# ema
train.model_ema.enabled=True
train.model_ema.decay=0.999

model.num_queries = 2000

model.transformer.encoder.use_checkpoint=False 
model.transformer.decoder.use_checkpoint=False

train.model_ema.use_ema_weights_for_eval_only=True