from yacs.config import CfgNode as CN


_C = CN()

_C.OUTPUT_DIR = ""
_C.DEVICE = "cuda"

_C.DATA = CN()
_C.DATA.DATASET = ""
_C.DATA.ROOT = ""
_C.DATA.TRAIN_FILENAME = ""
_C.DATA.TEST_FILENAME = ""

_C.NOAUG = CN()
_C.NOAUG.ENABLE = False
_C.NOAUG.SCALES = [1.0]

# input
_C.INPUT = CN()
_C.INPUT.SCALE_LOWER = 0.75
_C.INPUT.SCALE_UPPER = 1.33
_C.INPUT.PIXEL_MEAN = [0.485, 0.456, 0.406]
_C.INPUT.PIXEL_STD = [0.229, 0.224, 0.225]
_C.INPUT.TRAIN_IMG_SIZE = (512, 512)
_C.INPUT.TEST_IMG_SIZE = (512, 512)
_C.INPUT.NOAUG_IMG_SIZE = (512, 512)

# dataloder
_C.DATALOADER = CN()
_C.DATALOADER.NUM_WORKERS = 8
_C.DATALOADER.SAMPLER = "random"

# model
_C.MODEL = CN()
_C.MODEL.PRETRAIN = "imagenet"

_C.MODEL.BACKBONE = CN()
_C.MODEL.BACKBONE.NAME = "resnet101"

# pool
_C.MODEL.POOL = CN()
_C.MODEL.POOL.NAME = "gem"
_C.MODEL.POOL.OUTPUT_SIZE = 1

_C.MODEL.POOL.GMP = CN()
_C.MODEL.POOL.GMP.P = 3

# head
_C.MODEL.HEAD = CN()
_C.MODEL.HEAD.ENABLE = True
_C.MODEL.HEAD.NAME = "norm_linear_norm"
_C.MODEL.HEAD.IN_DIM = 2048
_C.MODEL.HEAD.OUT_DIM = 2048

# loss
_C.LOSS = CN()
_C.LOSS.NAME = "contrastive_loss"
_C.LOSS.CONTRASTIVE_LOSS = CN()
_C.LOSS.CONTRASTIVE_LOSS.MARGIN = 0.5
_C.LOSS.CONTRASTIVE_LOSS.POS_MARGIN = 0.0


_C.SOLVER = CN()

# optimizer
_C.SOLVER.NAME = "sgd"
_C.SOLVER.BASE_LR = 1e-3
_C.SOLVER.BIAS_LR_FACTOR = 1.0
_C.SOLVER.WEIGHT_DECAY = 1e-6
_C.SOLVER.WEIGHT_DECAY_BIAS = 0.0
_C.SOLVER.LOSS_BASE_LR = 1e-3
_C.SOLVER.LOSS_BIAS_LR_FACTOR = 1.0
_C.SOLVER.LOSS_WEIGHT_DECAY = 1e-6
_C.SOLVER.LOSS_WEIGHT_DECAY_BIAS = 0.0
_C.SOLVER.MOMENTUM = 0.9
_C.SOLVER.GRAD_MAX_NORM = 9999999

# scheduler
_C.SOLVER.SCHEDULER_NAME = "WarmupMultiStepLR"
_C.SOLVER.STEPS = [50000, 75000]
_C.SOLVER.MAX_STEP = 80000
_C.SOLVER.GAMMA = 0.1
_C.SOLVER.WARMUP_FACTOR = 1.0 / 3
_C.SOLVER.WARMUP_STEP = 0
_C.SOLVER.WARMUP_METHOD = "linear"

# training params
_C.SOLVER.IMS_PER_BATCH = 36
_C.SOLVER.CHECKPOINT_PERIOD = 4000

# test
_C.TEST = CN()
_C.TEST.TEST_PERIOD = 4000
_C.TEST.IMS_PER_BATCH = 128
_C.TEST.SCALES = [1, (1/2)**0.5, 1/2]


# xbm
_C.XBM = CN()
_C.XBM.VERSION = "v1"
_C.XBM.ENABLE = False
_C.XBM.POTENTIAL_POSITIVE = 50
_C.XBM.SIZE = 1223195 # dataset size
_C.XBM.FEATURE_DIM = 2048
_C.XBM.START_STEP = 2000
_C.XBM.GET_STEP = 4000
_C.XBM.GET_SIZE = 100000
_C.XBM.MEMORY_LOSS_ONLY = False
_C.XBM.MEMORY_LOSS_WEIGHT = 1.0
_C.XBM.USE_NEGATIVE = True

# for v2
_C.XBM.POSITIVE_THRESH = 0.65
_C.XBM.NEGATIVE_THRESH = 0.55
_C.XBM.CANDIDATE_MODE = ""

# neighbor
_C.NEIGHBOR = CN()
_C.NEIGHBOR.SRC = ""
_C.NEIGHBOR.FEATURE = ""
_C.NEIGHBOR.K = 1

_C.ATTENTION = CN()
_C.ATTENTION.ENABLE = False
_C.ATTENTION.NAME = "spaital_attention"
_C.ATTENTION.REDUCE_DIM = 128
_C.ATTENTION.REDUCE_RELU = False
_C.ATTENTION.FUNC = "sigmoid"

_C.SELECTION = CN()
_C.SELECTION.ENABLE = False
_C.SELECTION.USE_NOAUG_SIMILARITY = True
_C.SELECTION.THRESH = 0.65
_C.SELECTION.ALWAYS_TRUE = 0
_C.SELECTION.TYPE = "fixed"
