# --------------------------------------------------------
# Sequence Level Semantics Aggregation
# Copyright (c) 2016 by Contributors
# Copyright (c) 2017 Microsoft
# Copyright (c) 2019 by Contributors
# Licensed under The Apache-2.0 License [see LICENSE for details]
# Modified by Yuqing Zhu, Shuhao Fu, Xizhou Zhu, Yuwen Xiong, Bin Xiao
# Modified by Haiping Wu
# --------------------------------------------------------

import yaml
import numpy as np
from easydict import EasyDict as edict

config = edict()

config.MXNET_VERSION = ''
config.output_path = ''
config.log_path = './logs'
config.symbol = ''
config.gpus = ''
config.CLASS_AGNOSTIC = True
config.SCALES = [(600, 1000)]  # first is scale (the shorter side); second is max size

config.use_tensorboard = False

config.SYNC_BN = False
config.SYNC_BN_WD = False

# default training
config.default = edict()
config.default.frequent = 20
config.default.kvstore = 'device'

# network related params
config.network = edict()
config.network.pretrained = ''
config.network.pretrained_flow = ''
config.network.pretrained_fgfa = ''
config.network.pretrained_rfcn = ''
config.network.pretrained_epoch = 0
config.network.PIXEL_MEANS = np.array([0, 0, 0])
config.network.IMAGE_STRIDE = 0
config.network.RPN_FEAT_STRIDE = 16
config.network.RCNN_FEAT_STRIDE = 16
config.network.FIXED_PARAMS = ['gamma', 'beta']
config.network.ANCHOR_SCALES = (8, 16, 32)
config.network.ANCHOR_RATIOS = (0.5, 1, 2)
config.network.NORMALIZE_RPN = True
config.network.ANCHOR_MEANS = (0.0, 0.0, 0.0, 0.0)
config.network.ANCHOR_STDS = (0.1, 0.1, 0.4, 0.4)
config.network.NUM_ANCHORS = len(config.network.ANCHOR_SCALES) * len(config.network.ANCHOR_RATIOS)
config.network.FGFA_FEAT_DIM = 1024 + 2048 # 1024 for feature network, 2048 for embedding network

config.network.dtype = 'float32'


# dataset related params
config.dataset = edict()
config.dataset.dataset = 'ImageNetVID'
config.dataset.image_set = 'DET_train_30classes+VID_train_15frames'
config.dataset.test_image_set = 'VID_val_videos'
config.dataset.root_path = './data'
config.dataset.dataset_path = './data/ILSVRC2015'
config.dataset.motion_iou_path = './lib/dataset/imagenet_vid_groundtruth_motion_iou.mat'
config.dataset.enable_detailed_eval = True
config.dataset.NUM_CLASSES = 31


config.TRAIN = edict()

config.TRAIN.single_frame = False
config.TRAIN.lr = 0
config.TRAIN.lr_step = ''
config.TRAIN.lr_factor = 0.1
config.TRAIN.warmup = False
config.TRAIN.warmup_lr = 0
config.TRAIN.warmup_step = 0
config.TRAIN.momentum = 0.9
config.TRAIN.wd = 0.0005
config.TRAIN.begin_epoch = 0
config.TRAIN.end_epoch = 0
config.TRAIN.model_prefix = ''


# whether resume training
config.TRAIN.RESUME = False
# whether flip image
config.TRAIN.FLIP = True
# whether shuffle image
config.TRAIN.SHUFFLE = True
# whether use OHEM
config.TRAIN.ENABLE_OHEM = False
# size of images for each device, 2 for rcnn, 1 for rpn and e2e
config.TRAIN.BATCH_IMAGES = 2
# e2e changes behavior of anchor loader and metric
config.TRAIN.END2END = False
# group images with similar aspect ratio
config.TRAIN.ASPECT_GROUPING = True

# R-CNN
# rcnn rois batch size
config.TRAIN.BATCH_ROIS = 128
config.TRAIN.BATCH_ROIS_OHEM = 128
# rcnn rois sampling params
config.TRAIN.FG_FRACTION = 0.25
config.TRAIN.FG_THRESH = 0.5
config.TRAIN.BG_THRESH_HI = 0.5
config.TRAIN.BG_THRESH_LO = 0.0
# rcnn bounding box regression params
config.TRAIN.BBOX_REGRESSION_THRESH = 0.5
config.TRAIN.BBOX_WEIGHTS = np.array([1.0, 1.0, 1.0, 1.0])

# RPN anchor loader
# rpn anchors batch size
config.TRAIN.RPN_BATCH_SIZE = 256
# rpn anchors sampling params
config.TRAIN.RPN_FG_FRACTION = 0.5
config.TRAIN.RPN_POSITIVE_OVERLAP = 0.7
config.TRAIN.RPN_NEGATIVE_OVERLAP = 0.3
config.TRAIN.RPN_CLOBBER_POSITIVES = False
# rpn bounding box regression params
config.TRAIN.RPN_BBOX_WEIGHTS = (1.0, 1.0, 1.0, 1.0)
config.TRAIN.RPN_POSITIVE_WEIGHT = -1.0

# used for end2end training
# RPN proposal
config.TRAIN.CXX_PROPOSAL = True
config.TRAIN.RPN_NMS_THRESH = 0.7
config.TRAIN.RPN_PRE_NMS_TOP_N = 12000
config.TRAIN.RPN_POST_NMS_TOP_N = 2000
config.TRAIN.RPN_MIN_SIZE = config.network.RPN_FEAT_STRIDE
# approximate bounding box regression
config.TRAIN.BBOX_NORMALIZATION_PRECOMPUTED = True
config.TRAIN.BBOX_MEANS = (0.0, 0.0, 0.0, 0.0)
config.TRAIN.BBOX_STDS = (0.1, 0.1, 0.2, 0.2)

# FGFA, trained image sampled from [min_offset, max_offset]
config.TRAIN.MIN_OFFSET = -9
config.TRAIN.MAX_OFFSET = 9

config.TRAIN.append_gt = True

config.TRAIN.output_sort = False

config.TRAIN.train_frames = 3

config.TRAIN.loss_frames = 1

# use random flip augmentation for conditional frames other than be the same with the training frame
config.TRAIN.random_flip_condition_frames = False

config.TRAIN.DET_cls_sample = False
config.TRAIN.DET_negative_sample = False

config.TRAIN.loss_nonlocal_s = 1.0
config.TRAIN.loss_single_s = 0.0

# data augmentation
config.TRAIN.data_aug = False
config.TRAIN.COLOR_FACTOR = 0.0
config.TRAIN.ROT_FACTOR = 0.0
config.TRAIN.TRANS_FACTOR = 0.0
config.TRAIN.SCALE_FACTOR = 0.0
config.TRAIN.ssd_aug = False
config.TRAIN.noncur_aug = False
config.TRAIN.ssd_expand_scale = 2.0
config.TRAIN.ssd_crop_pert = 0.3
config.TRAIN.ssd_color = False
config.TRAIN.ssd_no_iou_limit = False

config.TEST = edict()

# R-CNN testing
# use rpn to generate proposal
config.TEST.HAS_RPN = False
# size of images for each device
config.TEST.BATCH_IMAGES = 1

# RPN proposal
config.TEST.CXX_PROPOSAL = True
config.TEST.RPN_NMS_THRESH = 0.7
config.TEST.RPN_PRE_NMS_TOP_N = 6000
config.TEST.RPN_POST_NMS_TOP_N = 300
config.TEST.RPN_MIN_SIZE = config.network.RPN_FEAT_STRIDE

# RCNN nms
config.TEST.NMS = 0.3
config.TEST.max_per_image = 300

#
config.TEST.KEY_FRAME_INTERVAL = 9
config.TEST.SEQ_NMS = False

# Test Model Epoch
config.TEST.test_epoch = 0

config.TEST.sample_stride = 1

config.TEST.video_shuffle = False

# whether performance feature aggregation if only one frame
config.TEST.single_frame_trans = False

config.RELATION = edict()
config.RELATION.conv_g = False
config.RELATION.inter_filter = 1024
config.RELATION.conv_z = True
config.RELATION.non_cur_space = False
config.RELATION.feature_wise_scale = False
config.RELATION.se_reduce = False
config.RELATION.feature_wise_scale_add = False
config.RELATION.feature_wise_scale_act = True
config.RELATION.feature_wise_add_out = False
config.RELATION.feature_wise_residual = False
config.RELATION.use_conv3_feat = False
config.RELATION.use_sim_label = False
config.RELATION.sim_sample_per_image = -1
config.RELATION.sample_ratio = -1
config.RELATION.ignore_bg_positive = 0
config.RELATION.relation_output_cur_only = False


def update_config(config_file):
    exp_config = None
    with open(config_file) as f:
        exp_config = edict(yaml.load(f))
        for k, v in exp_config.items():
            if k in config:
                if isinstance(v, dict):
                    if k == 'TRAIN':
                        if 'BBOX_WEIGHTS' in v:
                            v['BBOX_WEIGHTS'] = np.array(v['BBOX_WEIGHTS'])
                    elif k == 'network':
                        if 'PIXEL_MEANS' in v:
                            v['PIXEL_MEANS'] = np.array(v['PIXEL_MEANS'])
                    for vk, vv in v.items():
                        config[k][vk] = vv
                else:
                    if k == 'SCALES':
                        config[k][0] = (tuple(v))
                    else:
                        config[k] = v
            else:
                raise ValueError("key must exist in config.py")
