"""
This is the default experiment configuration for virat dataset. These values should
not be altered instead write a config file (in yaml) to override the default option.

Check sample_exp.yaml for a sample configuration file
"""
from easydict import EasyDict as edict
import os
import yaml
import numpy as np

experiment_config = edict()
# Virat data parameters
experiment_config.gpu                                  = 3
experiment_config.data                                 = edict()
experiment_config.data.frame_roots                     = ['/data/diva/v1-frames']
experiment_config.data.flow_roots                      = ['/data/diva/v1-flow-brox']
experiment_config.data.train_annotation_dirs           = ['/mnt/usb3/virat/nist-json/eval-1.B/train', '/mnt/usb3/virat/nist-json/eval-1.B/validate']
experiment_config.data.test_annotation_dirs            = ['/mnt/usb3/virat/nist-json/eval-1.B/test']
experiment_config.data.class_index                     = '/mnt/usb3/virat/nist-json/eval-1.B/test/activity-index.json'
experiment_config.data.save_directory                  = '/data/diva/baselines/activity_localization/act_detector/virat'
experiment_config.data.is_flow                         = False
experiment_config.data.save_prefix                     = 'virat'
experiment_config.data.img_extension                   = 'png'
experiment_config.data.num_frames                      = 6
    
# Training parameters
experiment_config.train                                = edict()
experiment_config.train.kpf_mode                       = False
experiment_config.train.json_mode                      = False
experiment_config.train.imgsize                        = 300
experiment_config.train.image_batch                    = 32
experiment_config.train.model_dir                      = '/data/diva/baselines/activity_localization/act_detector/virat'
experiment_config.train.pretrained_model_dir           = '/data/diva/baselines/activity_localization/act_detector/virat'
experiment_config.train.use_batchnorm                  = False
experiment_config.train.lr_mult                        = 1
experiment_config.train.batchnorm_lr                   = 0.0004
experiment_config.train.base_lr                        = 0.00004

# Solver parameters
experiment_config.train.solver                         = edict()
experiment_config.train.solver.niter                   = 600000
experiment_config.train.solver.lr_step                 = [400000, 550000]
experiment_config.train.solver.snapshot                = 10000
experiment_config.train.solver.display                 = 10
experiment_config.train.solver.average_loss            = 10
experiment_config.train.solver.snapshot_after_train    = True
experiment_config.train.solver.iter_size               = 1
experiment_config.train.solver.freeze_layers           = []
experiment_config.train.solver.weight_decay            = 0.0005
# Mbox parameters
experiment_config.train.mbox                           = edict()
experiment_config.train.mbox.share_location            = True
experiment_config.train.mbox.background_label_id       = 0
experiment_config.train.mbox.train_on_diff_gt          = True
experiment_config.train.mbox.ignore_cross_boundary_bbox = False
experiment_config.train.mbox.neg_pos_ratio             = 3.
experiment_config.train.mbox.overlap_threshold         = 0.5
experiment_config.train.mbox.use_prior_for_matching    = True
experiment_config.train.mbox.negative_overlap          = 0.5

# Priors parameters
experiment_config.train.priors                          = edict()
experiment_config.train.priors.src_layers              = ['conv4_3', 'fc7', 'conv6_2', 'conv7_2', 'conv8_2', 'conv9_2']
experiment_config.train.priors.min_ratio               = 20
experiment_config.train.priors.max_ratio               = 90
experiment_config.train.priors.aspect_ratios           = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
experiment_config.train.priors.normalizations          = [20, -1, -1, -1, -1, -1]
experiment_config.train.priors.steps                   = [8, 16, 32, 64, 100, 300] 
experiment_config.train.priors.flip                    = True
experiment_config.train.priors.clip                    = False

# Detection parameters
experiment_config.train.detection                      = edict()
experiment_config.train.detection.nms_threshold        = 0.45
experiment_config.train.detection.top_k                = 400
experiment_config.train.detection.keep_top_k           = 200
experiment_config.train.detection.confidence_threshold = 0.01

experiment_config.test = edict()
experiment_config.test.kpf_mode                        = False
experiment_config.test.json_mode                       = False
experiment_config.test.channel_mean                    = np.array([[[103, 117, 123]]], 
                                                                    dtype=np.float32)
experiment_config.test.number_flow                     = 5
experiment_config.test.batch_size                      = 5
experiment_config.test.detection_dir                   = "/data/diva/baselines/darknet-inference" 
experiment_config.test.cache_dir                       = "/data/diva/baselines/activity_localization/act_detector/models/ACT-detector/virat/eval_1.B/output_cache"

def merge_config(yaml_config, exp_config):
    if type(yaml_config) is not edict:
        return
    for key, value in yaml_config.iteritems():
        if not exp_config.has_key(key):
            raise KeyError('{} is not valid experiment key'.format(key))
        exp_type = type(exp_config[key])
        if exp_type is not type(yaml_config[key]):
            raise ValueError("Type mismatch ({} and {})".
                    format(type(exp_type), yaml_config[key]))
        elif type(value) is edict:
            merge_config(yaml_config[key], exp_config[key])
        else:
            exp_config[key] = value
            

def expcfg_from_file(filename):
    #Load a config file and merge it into the default options.
    with open(filename, 'r') as f:
        yaml_cfg = edict(yaml.load(f))
    merge_config(yaml_cfg, experiment_config)
