#!/usr/bin/python
# -*- coding:utf8 -*-
"""
    Author: Haoming Chen
    E-mail: chenhaomingbob@163.com
    Time: 2020/09/27
    Description:
"""
from .my_custom import CfgNode

_C = CfgNode()
_C.ROOT_DIR = ''
_C.EXPERIMENT_NAME = ''
_C.OUTPUT_DIR = ''
_C.SAVE_HEATMAPS = False
_C.LOAD_HEATMAPS = False
_C.SAVE_PREDS = False
_C.PREDS_SFX = ''
_C.LOAD_PREDS = False
_C.SAVE_OFFSETS = False
_C.LOG_DIR = ''
_C.DATA_DIR = ''
_C.MODEL_DIR = ''
_C.GPUS = (0,)
_C.WORKERS = 8
_C.PRINT_FREQ = 20
_C.PIN_MEMORY = True
_C.RANK = 0
_C.DISTANCE = 2
_C.NUMBER_SUP = 2
_C.CORE_FUNCTION = ""

_C.PATH_ADD_DESCRIPTIONS = CfgNode()
_C.PATH_ADD_DESCRIPTIONS.TRAIN = True
_C.PATH_ADD_DESCRIPTIONS.MODEL = True
_C.PATH_ADD_DESCRIPTIONS.DATASET = True
_C.PATH_ADD_DESCRIPTIONS.LOSS = True

# Cudnn related params
_C.CUDNN = CfgNode()
_C.CUDNN.BENCHMARK = True
_C.CUDNN.DETERMINISTIC = False
_C.CUDNN.ENABLED = True

# common params for NETWORK
_C.MODEL = CfgNode()
_C.MODEL.NAME = 'pose_hrnet'
_C.MODEL.INIT_WEIGHTS = True
_C.MODEL.FREEZE_WEIGHTS = False
_C.MODEL.FREEZE_PredNet_WEIGHTS = True
_C.MODEL.PRETRAINED = ''
_C.MODEL.BACKBONE_PRETRAINED = ''
_C.MODEL.NUM_JOINTS = 17
_C.MODEL.TARGET_TYPE = 'gaussian'
_C.MODEL.IMAGE_SIZE = [256, 256]  # width * height, ex: 192 * 256
_C.MODEL.HEATMAP_SIZE = [64, 64]  # width * height, ex: 24 * 32
_C.MODEL.SIGMA = 2
_C.MODEL.EXTRA = CfgNode(new_allowed=True)
_C.MODEL.CYCLE_CONSISTENCY_FINETUNE = False
_C.MODEL.DEFORAM_CONV_VERSION = 1
_C.MODEL.DEFORMABLE_CONV = CfgNode(new_allowed=True)
_C.MODEL.USE_RECTIFIER = True
_C.MODEL.USE_MARGIN = True
_C.MODEL.USE_GROUP = True
_C.MODEL.HIGH_RESOLUTION = False
_C.MODEL.FREEZE_HRNET_WEIGHTS = False
_C.MODEL.MPII_PRETRAINED = False
_C.MODEL.USE_WARPING_TRAIN = True
_C.MODEL.USE_WARPING_TEST = True
_C.MODEL.WARPING_REVERSE = False
_C.MODEL.USE_GT_INPUT_TEST = False
_C.MODEL.USE_GT_INPUT_TRAIN = False
_C.MODEL.ITER = 30000
_C.MODEL.EVALUATE = True
_C.MODEL.DILATION_EXP = 0
_C.MODEL.VISUALIZE_OFFSETS = False
_C.MODEL.USE_PIXEL_LEVEL_OFFSET = True
_C.MODEL.USE_PRF = True
_C.MODEL.PRF_BASICBLOCK_NUM = 10
_C.MODEL.PRF_INNER_CH = 12
_C.MODEL.USE_PTM = True
_C.MODEL.PTM_BASICBLOCK_NUM = 10
_C.MODEL.PTM_INNER_CH = 12
_C.MODEL.PRF_PTM_COMBINE_INNER_CH = 10
_C.MODEL.PRF_PTM_COMBINE_BASICBLOCK_NUM = 10
_C.MODEL.USE_PCN = True
_C.MODEL.TEMPORAL_INTERPOLATION = False
_C.MODEL.BACKBONE_PRECOMPUTE = False
_C.MODEL.WITH_DCPOSE = False
_C.MODEL.WARP_LEVEL = "Image"  # decrepit
_C.MODEL.LOCAL_WARP_LEVEL = "Image"  # decrepit
_C.MODEL.GLOBAL_WARP_LEVEL = "Image"  # decrepit

_C.MODEL.GLOBAL_WARP = CfgNode()
_C.MODEL.GLOBAL_WARP.LEVEL = "Patch"  # decrepit
_C.MODEL.GLOBAL_WARP.PATCH_WINDOW_SIZE = (4, 3)  # decrepit
_C.MODEL.GLOBAL_WARP.PATCH_WINDOW_STRIDE = (4, 3)  # decrepit

# 特征图的全局对齐设置
_C.MODEL.GLOBAL_WARP.FEATMAP = CfgNode()
_C.MODEL.GLOBAL_WARP.FEATMAP.LEVEL = "Image"
_C.MODEL.GLOBAL_WARP.FEATMAP.PATCH_WINDOW_SIZE = (96, 72)
_C.MODEL.GLOBAL_WARP.FEATMAP.PATCH_WINDOW_STRIDE = (96, 72)

# 热图的全局对齐设置
_C.MODEL.GLOBAL_WARP.HEATMAP = CfgNode()
_C.MODEL.GLOBAL_WARP.HEATMAP.LEVEL = "Image"
_C.MODEL.GLOBAL_WARP.HEATMAP.PATCH_WINDOW_SIZE = (96, 72)
_C.MODEL.GLOBAL_WARP.HEATMAP.PATCH_WINDOW_STRIDE = (96, 72)

_C.MODEL.LOCAL_WARP = CfgNode()
_C.MODEL.LOCAL_WARP.LEVEL = "Patch"  # decrepit
_C.MODEL.LOCAL_WARP.PATCH_WINDOW_SIZE = (12, 9)  # decrepit
_C.MODEL.LOCAL_WARP.PATCH_WINDOW_STRIDE = (12, 9)  # decrepit

# 特征图的局部对齐设置
_C.MODEL.LOCAL_WARP.FEATMAP = CfgNode()
_C.MODEL.LOCAL_WARP.FEATMAP.LEVEL = "Patch"
_C.MODEL.LOCAL_WARP.FEATMAP.PATCH_WINDOW_SIZE = (12, 9)
_C.MODEL.LOCAL_WARP.FEATMAP.PATCH_WINDOW_STRIDE = (12, 9)

# 热图的局部对齐设置
_C.MODEL.LOCAL_WARP.HEATMAP = CfgNode()
_C.MODEL.LOCAL_WARP.HEATMAP.LEVEL = "Image"
_C.MODEL.LOCAL_WARP.HEATMAP.PATCH_WINDOW_SIZE = (96, 72)
_C.MODEL.LOCAL_WARP.HEATMAP.PATCH_WINDOW_STRIDE = (96, 72)

#### LOSS ####
_C.LOSS = CfgNode()

_C.LOSS.AVG_LOSS = False
_C.LOSS.GRAD_MAX_NORM = 0.02

##
_C.LOSS.MI_SUMMATION_WEIGHT = 1.0  # \alpha
_C.LOSS.COMPLEMENTARY = CfgNode()
_C.LOSS.COMPLEMENTARY.USE = False
_C.LOSS.COMPLEMENTARY.WEIGHT = 0.1

_C.LOSS.VANISHING = CfgNode()
_C.LOSS.VANISHING.USE = False
_C.LOSS.VANISHING.WEIGHT = 1.0

_C.LOSS.IMAGE_RECON = CfgNode()
_C.LOSS.IMAGE_RECON.USE = False
_C.LOSS.IMAGE_RECON.WEIGHT = 0.5

_C.LOSS.CONSISTENCY = CfgNode()
_C.LOSS.CONSISTENCY.USE = True
_C.LOSS.CONSISTENCY.WEIGHT = 1.0

# _C.LOSS.MI_TERM_PAIR_1 = 0.1
# _C.LOSS.MI_TERM_PAIR_2 = 1.0
# _C.LOSS.MI_TERM_PAIR_3 = 1.0

## MSE
_C.LOSS.FM_GLOBAL_LOCAL = CfgNode()
_C.LOSS.FM_GLOBAL_LOCAL.USE = False
_C.LOSS.FM_GLOBAL_LOCAL.WEIGHT = 0.5

_C.LOSS.FM_GLOBAL_KF = CfgNode()
_C.LOSS.FM_GLOBAL_KF.USE = False
_C.LOSS.FM_GLOBAL_KF.WEIGHT = 0.5

_C.LOSS.HEATMAP_MSE = CfgNode()
_C.LOSS.HEATMAP_MSE.USE = True
_C.LOSS.HEATMAP_MSE.WEIGHT = 1.0
_C.LOSS.HEATMAP_MSE.DIVIDED_NUM_JOINTS = True

_C.LOSS.FEATMAP_MSE = CfgNode()
_C.LOSS.FEATMAP_MSE.USE = False
_C.LOSS.FEATMAP_MSE.WEIGHT = 0.5

_C.LOSS.LOCAL_HM = CfgNode()
_C.LOSS.LOCAL_HM.USE = True
_C.LOSS.LOCAL_HM.WEIGHT = 1.0

_C.LOSS.GLOBAL_HM = CfgNode()
_C.LOSS.GLOBAL_HM.USE = True
_C.LOSS.GLOBAL_HM.WEIGHT = 0.5

_C.LOSS.LOCAL_FEAT = CfgNode()
_C.LOSS.LOCAL_FEAT.USE = False
_C.LOSS.LOCAL_FEAT.WEIGHT = 0.8

_C.LOSS.GLOBAL_FEAT = CfgNode()
_C.LOSS.GLOBAL_FEAT.USE = False
_C.LOSS.GLOBAL_FEAT.WEIGHT = 0.2

_C.LOSS.GLOBAL_ALIGNMENT = CfgNode()
_C.LOSS.GLOBAL_ALIGNMENT.USE = False
_C.LOSS.GLOBAL_ALIGNMENT.WEIGHT = 0.01

_C.LOSS.LOCAL_ALIGNMENT = CfgNode()
_C.LOSS.LOCAL_ALIGNMENT.USE = False
_C.LOSS.LOCAL_ALIGNMENT.WEIGHT = 0.1

_C.LOSS.OFFSET_WARM_UP_EPOCH = 0
_C.LOSS.OFFSET = CfgNode()
_C.LOSS.OFFSET.USE = False
_C.LOSS.OFFSET.WEIGHT = 0.03

_C.LOSS.KL = CfgNode()
_C.LOSS.KL.USE = False
_C.LOSS.KL.WEIGHT = 0.01

_C.LOSS.BOUNDARY = CfgNode()
_C.LOSS.BOUNDARY.USE = False
_C.LOSS.BOUNDARY.WEIGHT = 0.01

_C.LOSS.DIVERSITY = CfgNode()
_C.LOSS.DIVERSITY.USE = False
_C.LOSS.DIVERSITY.WEIGHT = 0.01
# "["MSE", "NLL", "CrossEntropy", "KLDiv"]"
_C.LOSS.DIVERSITY.CRITERION = "MSE"

_C.LOSS.INTEGRAL_L1 = CfgNode()
_C.LOSS.INTEGRAL_L1.USE = False
_C.LOSS.INTEGRAL_L1.WEIGHT = 1.0

_C.LOSS.ALIGNED_FEAT = CfgNode()
_C.LOSS.ALIGNED_FEAT.USE = False
_C.LOSS.ALIGNED_FEAT.WEIGHT = 1.0

_C.LOSS.STRUCTURE_COSINE = CfgNode()
_C.LOSS.STRUCTURE_COSINE.USE = False
_C.LOSS.STRUCTURE_COSINE.WEIGHT = 1.0

_C.LOSS.OPTIMAL_TRANSPORT = CfgNode()
_C.LOSS.OPTIMAL_TRANSPORT.USE = False
_C.LOSS.OPTIMAL_TRANSPORT.EPSILON = 100  # [1, 10, 100]
_C.LOSS.OPTIMAL_TRANSPORT.N_ITER = 10  # [10 ** 4, 10 ** 3, 10 ** 2, 10]
_C.LOSS.OPTIMAL_TRANSPORT.WEIGHT = 1
# _C.LOSS.USE_HEATMAP_MSE = True
# _C.LOSS.USE_INTEGRAL_L1 = False
# _C.LOSS.USE_STRUCTURE_COSINE = False
# _C.LOSS.NAMES = ['MSELOSS']
# _C.LOSS.USE_OHKM = False
# _C.LOSS.USE_TARGET_WEIGHT = True
_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False
# _C.LOSS.USE_SOFTARGMAX = False

#### DATASET ####
_C.DATASET = CfgNode()
_C.DATASET.RANDOM_AUX_FRAME = True
_C.DATASET.ROOT = ''
_C.DATASET.NAME = ''
_C.DATASET.DATASET = 'mpii'
_C.DATASET.TRAIN_SET = 'train'
_C.DATASET.TEST_SET = 'test'
_C.DATASET.VAL_SET = 'val'
_C.DATASET.HYBRID_JOINTS_TYPE = ''
_C.DATASET.SELECT_DATA = False
_C.DATASET.TEST_ON_TRAIN = False
_C.DATASET.JSON_FILE = ''
_C.DATASET.JSON_DIR = ''
_C.DATASET.POSETRACK17_JSON_DIR = ''
_C.DATASET.POSETRACK18_JSON_DIR = ''
_C.DATASET.IMG_DIR = ''
_C.DATASET.POSETRACK17_IMG_DIR = ''
_C.DATASET.POSETRACK18_IMG_DIR = ''
_C.DATASET.IS_POSETRACK18 = False
_C.DATASET.COLOR_RGB = False
_C.DATASET.TEST_IMG_DIR = ''
_C.DATASET.POSETRACK17_TEST_IMG_DIR = ''
_C.DATASET.POSETRACK18_TEST_IMG_DIR = ''
_C.DATASET.INPUT_TYPE = ''
_C.DATASET.BBOX_ENLARGE_FACTOR = 1.0

_C.DATASET.USE_GLOBAL_REF = False
_C.DATASET.USE_LOCAL_REF = False
_C.DATASET.NUM_REF = 0
##
_C.DATASET.SPLIT_VERSION = 1

#### TRAIN ####
_C.TRAIN = CfgNode()  # cfg.Node
_C.TRAIN.SAVE_MODEL_PER_EPOCH = 2
_C.TRAIN.BATCH_SIZE_PER_GPU = 32
_C.TRAIN.SHUFFLE = True
_C.TRAIN.LOSS_ALPHA = 1.0
_C.TRAIN.LOSS_BETA = 1.0
_C.TRAIN.LOSS_GAMA = 1.0
_C.TRAIN.LR_FACTOR = 0.1
_C.TRAIN.LR_STEP = [90, 110]
_C.TRAIN.MILESTONES = [8, 12, 16]
_C.TRAIN.GAMMA = 0.99
_C.TRAIN.LR = 0.001
_C.TRAIN.STSN_LR = 0.001
_C.TRAIN.OPTIMIZER = 'adam'
_C.TRAIN.MOMENTUM = 0.9
_C.TRAIN.WD = 0.0001
_C.TRAIN.NESTEROV = False
_C.TRAIN.GAMMA1 = 0.99
_C.TRAIN.GAMMA2 = 0.0
_C.TRAIN.BEGIN_EPOCH = 0
_C.TRAIN.END_EPOCH = 140
_C.TRAIN.AUTO_RESUME = False
_C.TRAIN.FLIP = True
_C.TRAIN.SCALE_FACTOR = 0.25
_C.TRAIN.ROT_FACTOR = 30
_C.TRAIN.PROB_HALF_BODY = 0.0
_C.TRAIN.NUM_JOINTS_HALF_BODY = 8
_C.TRAIN.LR_SCHEDULER = 'MultiStepLR'
##
_C.TRAIN.LR_SECOND_GROUP = [None]
_C.TRAIN.LR_SECOND_GROUP_VALUE = 1e-6
_C.TRAIN.RANDOM_SAMPLE_IN_ENTIRE_TRACK_SEQUENCE = False
_C.TRAIN.SAMPLE_MAX_DISTANCE = 1
_C.TRAIN.BIDIRECTIONAL_SUPERVISION = False
_C.TRAIN.TRACK_SEQ = True
_C.TRAIN.TRAIN_GT_HEATMAPS_TRANSFORM = True
_C.TRAIN.TRAIN_AGG = False

#### VAL ####
_C.VAL = CfgNode()
_C.VAL.BATCH_SIZE_PER_GPU = 1
_C.VAL.MODEL_FILE = ''
_C.VAL.ANNOT_DIR = ''
_C.VAL.COCO_BBOX_FILE = ''
_C.VAL.USE_GT_BBOX = False
_C.VAL.FLIP_VAL = False
_C.VAL.BBOX_THRE = 1.0
# _C.VAL.BBOX_THRE = 1.0
_C.VAL.IMAGE_THRE = 0.1
_C.VAL.IN_VIS_THRE = 0.0
_C.VAL.NMS_THRE = 0.6
_C.VAL.OKS_THRE = 0.5
_C.VAL.SHIFT_HEATMAP = False
_C.VAL.SOFT_NMS = False
_C.VAL.POST_PROCESS = False
_C.VAL.FLIP = False
#### TEST ####
_C.TEST = CfgNode()
_C.TEST.BATCH_SIZE_PER_GPU = 1
_C.TEST.MODEL_FILE = ''
_C.TEST.ANNOT_DIR = ''
_C.TEST.COCO_BBOX_FILE = ''
_C.TEST.USE_GT_BBOX = False
_C.TEST.FLIP_TEST = False
_C.TEST.BBOX_THRE = 1.0
_C.TEST.IMAGE_THRE = 0.1
_C.TEST.IN_VIS_THRE = 0.0
_C.TEST.NMS_THRE = 0.6
_C.TEST.OKS_THRE = 0.5
_C.TEST.SHIFT_HEATMAP = False
_C.TEST.SOFT_NMS = False
_C.TEST.POST_PROCESS = False
_C.TEST.FLIP = False
### INFERENCE ###
_C.INFERENCE = CfgNode()
_C.INFERENCE.MODEL_FILE = ''
# debug
_C.DEBUG = CfgNode()
_C.DEBUG.VIS_SKELETON = False
_C.DEBUG.VIS_BBOX = False
_C.DEBUG.DEBUG = False
_C.DEBUG.SAVE_BATCH_IMAGES_GT = False
_C.DEBUG.SAVE_BATCH_IMAGES_PRED = False
_C.DEBUG.SAVE_HEATMAPS_GT = False
_C.DEBUG.SAVE_HEATMAPS_PRED = False
