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

from yacs.config import CfgNode as CN

"""default config for Quantization Block """
_QB = CN()

# OPERATION : {validate, train, collect_statistics}
_QB.OPERATION = 'validate'
_QB.OUTPUT_BASE_PATH = ''
_QB.INPUT_BASE_PATH = ''
_QB.DEVICE = 0
_QB.NUM_BITS = 8
_QB.EXP_NAME = ''
_QB.APPLY_FULL = False
_QB.DENSE_SPARSE = False
# DYNAMIC QUANTIZATION
_QB.DYNAMIC = CN()
_QB.DYNAMIC.APPLY = False
_QB.DYNAMIC.UP_QUANTILE = 1.0
_QB.DYNAMIC.LOW_QUANTILE = 1.0



# STATIC QUANTIZATION
_QB.STATIC = CN()
_QB.STATIC.APPLY = False
_QB.STATIC.UP_THRESHOLDS_PATH = ''
_QB.STATIC.LOW_THRESHOLDS_PATH = ''


# QPP QUANTIZATION
_QB.QPP = CN()
_QB.QPP.APPLY = False
_QB.QPP.UP_PREDICTORS_PATH = ''
_QB.QPP.LOW_PREDICTORS_PATH = ''





# COLLECT STATISTICS PARAMETERS
_QB.COLLECT = CN()
_QB.COLLECT.STATISTICS_PATH = ''
_QB.COLLECT.UP_THRESHOLDS_PATH = ''
_QB.COLLECT.LOW_THRESHOLDS_PATH = ''
_QB.COLLECT.APPLY = False
_QB.COLLECT.UP_QUANTILES = []
_QB.COLLECT.LOW_QUANTILES = []
_QB.COLLECT.MEAN = True
_QB.COLLECT.STD = True
_QB.COLLECT.MAX = True
_QB.COLLECT.MIN = True
_QB.COLLECT.ABS_MAX = True



def update_config_QB(cfg, qb_filename):
    cfg.defrost()
    cfg.merge_from_file(qb_filename)
    cfg.freeze()


if __name__ == '__main__':
    import sys

    with open(sys.argv[1], 'w') as f:
        print(_QB, file=f)