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

from yacs.config import CfgNode as CN

""" Default config for experiment set up"""
_EXPs = CN()

# OPERATION : {validate, train, collect_statistics}
_EXPs.OPERATION = ['validate']
_EXPs.EXP_NAMES = ['']
_EXPs.SOURCE = ''
_EXPs.OUTPUT_BASE_PATH = ''
_EXPs.INPUT_BASE_PATH = ''
_EXPs.LOG = 'log'
_EXPs.WRITERS = 'writers'
# one QB_CONFIG for every EXP
# INPUT_BASE_PATH +  './qb_configs/qb_config_name'
_EXPs.QB_CONFIGS = ['']
# one checkpoint for every EXP
# CHECKPOINTS : INPUT_BASE_PATH +  './checkpoints/checkpoint_name'
_EXPs.CHECKPOINTS = ['']
# one model for every EXP
# MODEL_CONFIGS : INPUT_BASE_PATH +  './model_configs/model_config_name'
# one device for every experiment
_EXPs.DEVICES = ['']
_EXPs.NN_TYPE = 'resnet18'
# DATASETS PATHS:
# TYPE OF DATASET
_EXPs.DATASET_TYPE = ''
_EXPs.TRAIN_DATA = ''
_EXPs.TEST_DATA = ''
# weights quantization parameters
_EXPs.ACT_NUM_BITS = [8]
_EXPs.DENSE_SPARSE = [False]
_EXPs.WEIGHT = CN()
_EXPs.WEIGHT.NUM_BITS = [8]
_EXPs.WEIGHT.Q_ALGORITHM = ['asymmetric']

# VALIDATION PARAMs
_EXPs.VALIDATE = CN()
_EXPs.VALIDATE.BATCH_SIZE_PER_GPU = 16
_EXPs.VALIDATE.SAVE_IMAGE_PATH = ['images']
# TRAIN PARAMs
_EXPs.TRAIN = CN()
_EXPs.TRAIN.OPTIMIZER = ['SGD']
_EXPs.TRAIN.MOMENTUM = [0.0]
_EXPs.TRAIN.NESTEROV = [False]
_EXPs.TRAIN.LR = [0.0001]
_EXPs.TRAIN.WD = [0.0005]
_EXPs.TRAIN.SAVE_CHECKPOINT_PATH = ['checkpoints']
_EXPs.TRAIN.BATCH_SIZE_PER_GPU = [16]
_EXPs.TRAIN.SHUFFLE = [True]
_EXPs.TRAIN.BEGIN_EPOCH = [0]
_EXPs.TRAIN.END_EPOCH = [30]
_EXPs.TRAIN.RESUME = [True]


# COLLECT STATISTICS
_EXPs.COLLECT = CN()
_EXPs.COLLECT.STATISTICS_PATH = ''

def update_config_EXP(cfg,args ):
    cfg.defrost()
    cfg.merge_from_file(args.cfg_exp)
    cfg.freeze()


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



