
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from yacs.config import CfgNode as CN


_C = CN()

_C.OUTPUT_DIR = ''
_C.LOG_DIR = ''
_C.DATA_DIR = ''
_C.GPUS = (0,)
_C.WORKERS = 4
_C.PRINT_FREQ = 20
_C.AUTO_RESUME = False
_C.PIN_MEMORY = True
_C.RANK = 0

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

# common params for NETWORK
_C.MODEL = CN()
_C.MODEL.NAME = 'pose_hrnet'
_C.MODEL.INIT_WEIGHTS = True
_C.MODEL.PRETRAINED = ''
_C.MODEL.NUM_JOINTS = 17
_C.MODEL.NUM_MAPS = 30
_C.MODEL.TAG_PER_JOINT = True
_C.MODEL.TARGET_TYPE = 'gaussian'
_C.MODEL.IS_IMM = False
_C.MODEL.IS_GT_GENERATION = False
_C.MODEL.PREDICT_UNSUP_PTS = True
_C.MODEL.MIX_GT = False
_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.CROP_BACKGROUND = False
_C.MODEL.CROP_BACKGROUND_THRESHOLD = 0.1
_C.MODEL.EXTRA = CN(new_allowed=True)

_C.LOSS = CN()
_C.LOSS.USE_OHKM = False
_C.LOSS.TOPK = 8
_C.LOSS.USE_TARGET_WEIGHT = True
_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False
_C.LOSS.UNSUP_LOSS_WEIGHT = 1.0
_C.LOSS.CROP = False
_C.LOSS.CROP_MARGIN = 5
_C.LOSS.NEW_VGG_LOSS = False
_C.LOSS.L2_LOSS = False
_C.LOSS.WEIGHT_SIN = False

# DATASET related params
_C.DATASET = CN()
_C.DATASET.ROOT = ''
_C.DATASET.DATASET = 'mpii'
_C.DATASET.TRAIN_SET = 'train'
_C.DATASET.TEST_SET = 'valid'
_C.DATASET.DATA_FORMAT = 'jpg'
_C.DATASET.HYBRID_JOINTS_TYPE = ''
_C.DATASET.SELECT_DATA = False
_C.DATASET.TRAIN_VIDEO_FRAC = 1.0
_C.DATASET.TRAIN_FRAME_FRAC = 1.0
_C.DATASET.LABEL_VIDEO_FRAC = 1.0
_C.DATASET.LABEL_FRAME_FRAC = 1.0

# training data augmentation
_C.DATASET.FLIP = True
_C.DATASET.SCALE_FACTOR = 0.25
_C.DATASET.ROT_FACTOR = 30
_C.DATASET.PROB_HALF_BODY = 0.0
_C.DATASET.NUM_JOINTS_HALF_BODY = 8
_C.DATASET.COLOR_RGB = False
_C.DATASET.COLOR_JITTER = 0.0
_C.DATASET.GRAY_SCALE = 0.0
_C.DATASET.GAUSSIAN_BLUR = 0.0
_C.DATASET.SELECTED_ACTIONS = ['baseball_pitch', 'clean_and_jerk', 'pullup', 'strum_guitar',
                                'baseball_swing', 'golf_swing', 'pushup', 'tennis_forehand',
                                'bench_press', 'jumping_jacks', 'situp', 'tennis_serve',
                                'bowl', 'jump_rope', 'squat']
_C.DATASET.BANDWIDTH = 1000000  # no bandwidth restriction

# train
_C.TRAIN = CN()

_C.TRAIN.LR_FACTOR = 0.1
_C.TRAIN.LR_STEP = [90, 110]
_C.TRAIN.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.MASK_AS_INPUT = False
_C.TRAIN.MASK_AS_REF = False

_C.TRAIN.RESUME = False
_C.TRAIN.CHECKPOINT = ''

_C.TRAIN.BATCH_SIZE_PER_GPU = 32
_C.TRAIN.SHUFFLE = True

_C.TRAIN.TTT_BATCHING_STRATEGY = 'fix_one'

_C.TRAIN.VALIDATE_EVERY = 1
_C.TRAIN.LIMIT_EPOCH_SIZE = -1

# testing
_C.TEST = CN()

# size of images for each device
_C.TEST.BATCH_SIZE_PER_GPU = 32
# Test Model Epoch
_C.TEST.FLIP_TEST = False
_C.TEST.POST_PROCESS = False
_C.TEST.SHIFT_HEATMAP = False

_C.TEST.USE_GT_BBOX = False

# ttt
# reset weight every n frames, <=0 means no reset
_C.TEST.RESET_EVERY = -1
# step on sample sample for n steps
_C.TEST.REPEAT_STEP = 1
# step only on first frame
_C.TEST.FIRST_FRAME_ONLY = False
# Strategy for Reference Frame
_C.TEST.TTT_ONLINE_REF = True
# Train on the whole video first
_C.TEST.TTT_RESET = -1
_C.TEST.EXTRA_VAL = 0
_C.TEST.TRAIN_FIRST = False
_C.TEST.TTT_STEPS = 200
_C.TEST.REPEAT_STEP = 1
_C.TEST.TTT_ONLINE = False
_C.TEST.SMOOTH = False
_C.TEST.DOWNSAMPLE = 1

# nms
_C.TEST.IMAGE_THRE = 0.1
_C.TEST.NMS_THRE = 0.6
_C.TEST.SOFT_NMS = False
_C.TEST.OKS_THRE = 0.5
_C.TEST.IN_VIS_THRE = 0.0
_C.TEST.COCO_BBOX_FILE = ''
_C.TEST.BBOX_THRE = 1.0
_C.TEST.MODEL_FILE = ''
_C.TEST.POSE_NET_FILE = ''

# debug
_C.DEBUG = CN()
_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
_C.DEBUG.SAVE_REF_IMAGES = False
_C.DEBUG.SAVE_RENDERED_IMAGES = False
_C.DEBUG.SAVE_UNSUP_PRED = False
_C.DEBUG.SAVE_HEATMAPS_UNSUP = False


def update_config(cfg, args):
    cfg.defrost()
    cfg.merge_from_file(args.cfg)
    cfg.merge_from_list(args.opts)

    if args.modelDir:
        cfg.OUTPUT_DIR = args.modelDir

    if args.logDir:
        cfg.LOG_DIR = args.logDir

    if args.dataDir:
        cfg.DATA_DIR = args.dataDir

    cfg.DATASET.ROOT = os.path.join(
        cfg.DATA_DIR, cfg.DATASET.ROOT
    )

    cfg.MODEL.PRETRAINED = os.path.join(
        cfg.DATA_DIR, cfg.MODEL.PRETRAINED
    )

    if cfg.TEST.MODEL_FILE:
        cfg.TEST.MODEL_FILE = os.path.join(
            cfg.DATA_DIR, cfg.TEST.MODEL_FILE
        )

    cfg.freeze()


if __name__ == '__main__':
    import sys
    with open(sys.argv[1], 'w') as f:
        print(_C, file=f)

