from yacs.config import CfgNode as CN

cfg = CN(new_allowed=True)

cfg.MODELS = CN()
cfg.MODELS.MASKS = False
cfg.MODELS.FROZEN_WEIGHT = None
cfg.MODELS.BACKBONE = 'resnet50'
cfg.MODELS.DILATION = False
cfg.MODELS.POSITION_EMBEDDING = 'sine'
cfg.MODELS.USE_STRUCTURE_TENSOR = True

cfg.MODELS.TRANSFORMER = CN()
cfg.MODELS.TRANSFORMER.ENC_LAYERS = 6
cfg.MODELS.TRANSFORMER.DEC_LAYERS = 6
cfg.MODELS.TRANSFORMER.DIM_FEEDFORWARD = 2048
cfg.MODELS.TRANSFORMER.HIDDEN_DIM = 256
cfg.MODELS.TRANSFORMER.DROPOUT = 0.1
cfg.MODELS.TRANSFORMER.NHEADS = 8
cfg.MODELS.TRANSFORMER.NUM_QUERIES = 3
cfg.MODELS.TRANSFORMER.PRE_NORM = False

cfg.LOSS = CN()
cfg.LOSS.LOSSES = ['zvp','fovy','hl','consistency']
cfg.LOSS.AUX_LOSS = True
cfg.LOSS.LINE_POS_ANGLE = 88.0
cfg.LOSS.LINE_NEG_ANGLE = 85.0
cfg.LOSS.WEIGHTS = CN(new_allowed=True)

cfg.DATASETS = CN()
cfg.DATASETS.INPUT_WIDTH = 512
cfg.DATASETS.INPUT_HEIGHT = 512
cfg.DATASETS.MIN_LINE_LENGTH = 10
cfg.DATASETS.NUM_INPUT_LINES = 512
cfg.DATASETS.NUM_INPUT_VERT_LINE = 256
cfg.DATASETS.VERT_LINE_ANGLE = 22.5 # degrees
cfg.DATASETS.RETURN_VERT_LINES = False

cfg.SOLVER = CN()
cfg.SOLVER.LR = 1e-4
cfg.SOLVER.LR_BACKBONE = 1e-4
cfg.SOLVER.BATCH_SIZE = 56
cfg.SOLVER.WEIGHT_DECAY = 1e-4
cfg.SOLVER.EPOCHS = 50
cfg.SOLVER.LR_DROP = 20
cfg.SOLVER.CLIP_MAX_NORM = 0.1

cfg.TEST = CN()
cfg.TEST.DISPLAY = True

cfg.OUTPUT_DIR = "outputs"
cfg.DEVICE = 'cuda'
cfg.SEED = 42
cfg.RESUME = ''
cfg.START_EPOCH = 0
cfg.VAL = False
cfg.NUM_WORKERS = 12
cfg.WORLD_SIZE = 1
cfg.DIST_URL = 'tcp://127.0.0.1:23456'