#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import yaml
from schema import Schema, And, Use, Optional, Or, Hook
from ..adaptor import FRAMEWORKS
from ..strategy import STRATEGIES
from ..objective import OBJECTIVES
from ..pruners import PRUNERS
from ..utils import logger
from ..version import __version__
import re
import copy
import itertools
from collections import OrderedDict
from .dotdict import DotDict, deep_set
import os, datetime

def constructor_register(cls):
    yaml_key = "!{}".format(cls.__name__)

    def constructor(loader, node):
        instance = cls.__new__(cls)
        yield instance

        state = loader.construct_mapping(node, deep=True)
        instance.__init__(**state)

    yaml.add_constructor(
        yaml_key,
        constructor,
        yaml.SafeLoader,
    )
    return cls

@constructor_register
class Pruner():
    def __init__(self, start_epoch=None, end_epoch=None, initial_sparsity=None,
                 target_sparsity=None, update_frequency=1,
                 method='per_tensor',
                 prune_type='basic_magnitude',##for pytorch pruning, these values should be None
                 start_step=None, end_step=None, update_frequency_on_step=None, prune_domain=None,
                 sparsity_decay_type=None, pattern="tile_pattern_1x1", names=None,
                 extra_excluded_names=None, parameters=None):
        self.start_epoch = start_epoch
        self.end_epoch = end_epoch
        self.update_frequency = update_frequency
        self.target_sparsity = target_sparsity
        self.initial_sparsity = initial_sparsity
        self.update_frequency = update_frequency
        self.start_step = start_step
        self.end_step = end_step
        self.update_frequency_on_step = update_frequency_on_step
        self.prune_domain = prune_domain
        self.sparsity_decay_type = sparsity_decay_type
        self.extra_excluded_names = extra_excluded_names
        self.pattern = pattern
        ## move this to experimental/pruning to support dynamic pruning
        # assert prune_type.replace('_', '') in [i.lower() for i in PRUNERS], \
        #                                  'now only support {}'.format(PRUNERS.keys())
        self.prune_type = prune_type
        self.method = method
        self.names= names
        self.parameters = parameters

# Schema library has different loading sequence priorities for different
# value types.
# To make sure the fields under dataloader.transform field of yaml file
# get loaded with written sequence, this workaround is used to convert
# None to {} in yaml load().
yaml.SafeLoader.add_constructor('tag:yaml.org,2002:null', lambda loader, node: {})
# Add python tuple support because best_configure.yaml may contain tuple
yaml.SafeLoader.add_constructor('tag:yaml.org,2002:python/tuple',
                                lambda loader, node: tuple(loader.construct_sequence(node)))

def _valid_accuracy_field(key, scope, error):
    assert bool(
        'relative' in scope['accuracy_criterion']) != bool(
        'absolute' in scope['accuracy_criterion'])

def _valid_prune_epoch(key, scope, error):
    if "start_epoch" in scope[key] and "end_epoch" in scope[key]:
        assert scope[key]["start_epoch"] <= scope[key]["end_epoch"]

def _valid_prune_sparsity(key, scope, error):
    if "initial_sparsity" in scope[key] and "target_sparsity" in scope[key]:
        assert scope[key]["initial_sparsity"] <= scope[key]["target_sparsity"]
    elif "initial_sparsity" in scope[key]:
        assert scope[key]["initial_sparsity"] >= 0
    elif "target_sparsity" in scope[key]:
        assert scope[key]["target_sparsity"] < 1

def _valid_multi_objectives(key, scope, error):
    if 'weight' in scope[key] and scope[key]['weight'] is not None:
        assert len(scope[key]['objective']) == len(scope[key]['weight'])

def _valid_multi_metrics(key, scope, error):
    if 'metric' in scope and 'multi_metrics' in scope:
        assert False

def _valid_metric_length(key, scope, error):
    metrics = [i for i in scope[key] if i != 'weight' and i != 'higher_is_better']
    if 'weight' in scope[key] and scope[key]['weight'] is not None:
        assert len(input_to_list_float(scope[key]['weight'])) == len(metrics)
    if 'higher_is_better' in scope[key] and scope[key]['higher_is_better'] is not None:
        assert len(input_to_list_bool(scope[key]['higher_is_better'])) == len(metrics)

# used for '123.68 116.78 103.94' style to float list
def input_to_list_float(data):
    if isinstance(data, str):
        return [float(s.strip()) for s in data.split()]

    if isinstance(data, float):
        return [data]

    assert isinstance(data, list)
    return [float(d) for d in data]

def input_to_list_bool(data):
    if isinstance(data, str):
        if ',' in data:
            return [s.strip() == 'True' for s in data.split(',')]
        else:
            return [s.strip() == 'True' for s in data.split()]

    if isinstance(data, bool):
        return [data]

    assert isinstance(data, list) and all([isinstance(i, bool) for i in data])
    return data

def input_int_to_float(data):
    if isinstance(data, str):
        # used for '123.68, 116.78, 103.94' style
        if ',' in data:
            data = data.split(',')
        # used for '123.68 116.78 103.94' style
        else:
            data = data.split()

        if len(data) == 1:
            return float(data[0].strip())
        else:
            return [float(s.strip()) for s in data]
    elif isinstance(data, list):
        return [float(s) for s in data]
    elif isinstance(data, int):
        return float(data)

def input_to_list_int(data):
    if isinstance(data, str):
        return [int(s.strip()) for s in data.split(',')]

    if isinstance(data, int):
        return [data]

    assert isinstance(data, list)
    return [int(d) for d in data]

def input_to_list(data):
    if isinstance(data, str):
        if ',' in data:
            return [s.strip() for s in data.split(',')]

        return [s.strip() for s in data.split()]

    if isinstance(data, int):
        return [data]

    assert isinstance(data, list)
    return data

def list_to_tuple(data):
    if isinstance(data, str):
        return tuple([int(s.strip()) for s in data.split(',')])

    elif isinstance(data, list):
        if isinstance(data[0], list):
            result = []
            for item in data:
                result.append(tuple([int(s) for s in item]))
            return result
        else:
            return tuple([int(s) for s in data])

def percent_to_float(data):
    if isinstance(data, str) and re.match(r'-?\d+(\.\d+)?%', data):
        data = float(data.strip('%')) / 100
    if isinstance(data, int):
        data = float(data)
    else:
        assert isinstance(data, float), 'This field should be float, int or percent string'
    return data

ops_schema = Schema({
    Optional('weight', default=None): {
        Optional('granularity'): And(
            list,
            lambda s: all(i in ['per_channel', 'per_tensor'] for i in s)),
        Optional('scheme'): And(
            list,
            # asym_float are only for PyTorch framework
            lambda s: all(i in ['asym', 'sym', 'asym_float'] for i in s)),
        Optional('dtype'): And(
            list,
            lambda s: all(i in ['int8', 'uint8', 'fp32', 'bf16', 'fp16'] for i in s)),
        Optional('algorithm'): And(
            list,
            lambda s: all(i in ['minmax'] for i in s)),
        Optional('bit'):  And(
            Or(float, list),
            Use(input_to_list_float),
            lambda s: all(0.0 < i <= 7.0 for i in s))
    },
    Optional('activation', default=None): {
        Optional('granularity'): And(
            list,
            lambda s: all(i in ['per_channel', 'per_tensor'] for i in s)),
        Optional('scheme'): And(
            list,
            lambda s: all(i in ['asym', 'sym'] for i in s)),
        Optional('dtype'): And(
            list,
            lambda s: all(i in ['int8', 'uint8', 'fp32', 'bf16', 'fp16'] for i in s)),
        # compute_dtypeis only for PyTorch framework
        Optional('compute_dtype', default=['uint8']): And(
            list,
            lambda s: all(i in ['int8', 'uint8', 'fp32', 'bf16', 'None'] for i in s)),
        # placeholder are only for PyTorch framework
        Optional('algorithm'): And(
            list,
            lambda s: all(i in ['minmax', 'kl', 'placeholder'] for i in s))
    }
})

graph_optimization_schema = Schema({

    Optional('precisions', default={'precisions': ['fp32']}): And(
        Or(str, list),
        Use(input_to_list),
        lambda s: all(i in [ 'fp32', 'bf16'] for i in s)),

    Optional('op_wise', default={'weight': {}, 'activation': {}}): {
        Optional('weight', default=None): {
            Optional('dtype', default=None): And(
                Or(str, list),
                Use(input_to_list),
                lambda s: all(i in ['fp32', 'bf16', 'fp16'] for i in s)),
        },
        Optional('activation', default=None): {
            Optional('dtype', default=None): And(
                Or(str, list),
                Use(input_to_list),
                lambda s: all(i in ['fp32', 'bf16', 'fp16'] for i in s)),
            }
    }
})

mixed_precision_schema = Schema({

    Optional('precisions', default={'precisions': ['fp32']}): And(
        Or(str, list),
        Use(input_to_list),
        lambda s: all(i in [ 'fp32', 'bf16', 'fp16'] for i in s)),

    Optional('op_wise', default={'weight': {}, 'activation': {}}): {
        Optional('weight', default=None): {
            Optional('dtype', default=None): And(
                Or(str, list),
                Use(input_to_list),
                lambda s: all(i in ['fp32', 'bf16', 'fp16'] for i in s)),
        },
        Optional('activation', default=None): {
            Optional('dtype', default=None): And(
                Or(str, list),
                Use(input_to_list),
                lambda s: all(i in ['fp32', 'bf16', 'fp16'] for i in s)),
            }
    }
})

model_conversion_schema = Schema({
    'source': And(str, lambda s: s.lower() == 'qat'),
    'destination': And(str, lambda s: s.lower() == 'default')
})

filter_schema = Schema({
    Optional('LabelBalance'): {
        'size': And(int, lambda s: s > 0)
    },
})

transform_schema = Schema({
    Optional('ResizeWithRatio'):{
        Optional('min_dim'): int,
        Optional('max_dim'): int,
        Optional('padding'): bool,
        Optional('constant_value'): int
    },
    Optional('CropToBoundingBox'): {
        'offset_height': int,
        'offset_width': int,
        'target_height': int,
        'target_width': int
    },
    Optional('Cast'): {
        Optional('dtype'): str
    },
    Optional('RandomResizedCrop'): {
        'size': Or(And(list, lambda s: all(isinstance(i, int) for i in s)),
                    And(int, lambda s: s > 0)),
        Optional('scale'): And(list, lambda s: all(isinstance(i, float) for i in s)),
        Optional('ratio'): And(list, lambda s: all(isinstance(i, float) for i in s)),
        Optional('interpolation'): And(
            str,
            lambda s: s in ['nearest', 'bilinear', 'bicubic']),
    },
    Optional('AlignImageChannel'): {
        Optional('dim'): int
    },
    Optional('ToNDArray'): Or({}, None),
    Optional('CropResize'): {
        'x': int,
        'y': int,
        'width': int,
        'height': int,
        'size': Or(And(list, lambda s: all(isinstance(i, int) for i in s)),
                    And(int, lambda s: s > 0)),
        Optional('interpolation'): And(
            str,
            lambda s: s in ['nearest', 'bilinear', 'bicubic']),
    },
    Optional('RandomHorizontalFlip'): Or({}, None),
    Optional('RandomVerticalFlip'): Or({}, None),
    Optional('ToTensor'): Or({}, None),
    Optional('ToPILImage'): Or({}, None),
    Optional('Normalize'): {
        Optional('mean'): And(list, lambda s: all(isinstance(i, float) for i in s)),
        Optional('std'): And(list, lambda s: all(isinstance(i, float) for i in s)),
        Optional('rescale'): list
    },
    Optional('Resize'): {
        'size': Or(And(list, lambda s: all(isinstance(i, int) for i in s)),
                    And(int, lambda s: s > 0)),
        Optional('interpolation'): And(
            str,
            lambda s: s in ['nearest', 'bilinear', 'bicubic']),
    },
    Optional('RandomCrop'): {
        'size': Or(And(list, lambda s: all(isinstance(i, int) for i in s)),
                    And(int, lambda s: s > 0))
    },
    Optional('Rescale'): Or({}, None),
    Optional('CenterCrop'): {
        'size': Or(And(list, lambda s: all(isinstance(i, int) for i in s)),
                    And(int, lambda s: s > 0))
    },
    Optional('PaddedCenterCrop'): {
        'size': Or(And(list, lambda s: all(isinstance(i, int) for i in s)),
                    And(int, lambda s: s > 0)),
        Optional('crop_padding'): And(int, lambda s: s > 0),
    },
    Optional('ToArray'): Or({}, None),
    Optional('QuantizedInput'): {
        Optional('dtype', default='int8'): And(str, lambda s: s in ['int8', 'uint8']),
        Optional('scale'): And(float, lambda s: s > 0),
    },
    Optional('Transpose'): {
        'perm': And(list, lambda s: all(isinstance(i, int) for i in s)),
    },
    # THIS API IS TO BE DEPRECATED!
    Optional('ParseDecodeImagenet'): Or({}, None),
    Optional('ParseDecodeCoco'): Or({}, None),
    Optional('ParseDecodeVoc'): Or({}, None),
    Optional('BilinearImagenet'): {
        'height': And(int, lambda s: s > 0),
        'width': And(int, lambda s: s > 0),
        Optional('central_fraction'): float,
        Optional('mean_value'): And(Or(str, list), Use(input_to_list_float)),
        Optional('scale'): float,
    },
    Optional('ResizeCropImagenet'): {
        'height': And(int, lambda s: s > 0),
        'width': And(int, lambda s: s > 0),
        Optional('random_crop'): bool,
        Optional('slice_crop'): bool,
        Optional('resize_side'): And(int, lambda s: s > 0),
        Optional('resize_method', default='bilinear'): \
            And(str, lambda s: s in ['bilinear', 'lanczos3', 'lanczos5',
                                     'bicubic', 'gaussian', 'nearest',
                                     'area', 'mitchellcubic']),
        Optional('random_flip_left_right'): bool,
        Optional('data_format', default='channels_last'): \
            And(str, lambda s: s in ['channels_first', 'channels_last']),
        Optional('subpixels', default='RGB'): \
            And(str, lambda s: s in ['BGR', 'RGB']),
        Optional('mean_value'): And(Or(str, list), Use(input_to_list_float)),
        Optional('scale'): float,
    },
    Optional('ResizeWithAspectRatio'):{
        'height': And(int, lambda s: s > 0),
        'width': And(int, lambda s: s > 0),
    },
    Optional('ParseDecodeImagenet'): Or({}, None),
    Optional('ToArray'): Or({}, None),
    Optional('QuantizedInput'): {
        Optional('dtype', default='int8'): And(str, lambda s: s in ['int8', 'uint8']),
        Optional('scale'): And(float, lambda s: s > 0),
    },
    Optional('Transpose'): {
        'perm': And(list, lambda s: all(isinstance(i, int) for i in s)),
    },
})

postprocess_schema = Schema({
    Optional('LabelShift'):  int,
    Optional('Collect'): {
        'length': int
    },
    Optional('SquadV1'): {
        'label_file': str,
        'vocab_file': str,
        Optional('do_lower_case', default='True'): bool,
        Optional('max_seq_length', default=384): int,
    },
    Optional('SquadV1ModelZoo'): {
        'label_file': str,
        'vocab_file': str,
        Optional('do_lower_case', default='True'): bool,
        Optional('max_seq_length', default=384): int,
    },
})

dataset_schema = Schema({
    Optional('CIFAR10'): {
        'root': str,
        Optional('train'): bool,
        Optional('download'): bool,
    },
    Optional('CIFAR100'): {
        'root': str,
        Optional('train'): bool,
        Optional('download'): bool,
    },
    Optional('MNIST'): {
        'root': str,
        Optional('train'): bool,
        Optional('download'): bool,
    },
    Optional('FashionMNIST'): {
        'root': str,
        Optional('train'): bool,
        Optional('download'): bool,
    },
    Optional('ImageFolder'): {
        'root': str,
    },
    Optional('TFRecordDataset'): {
        'root': str,
    },
    Optional('ImageRecord'): {
        'root': str,
    },
    Optional('dummy_v2'): {
        'input_shape': And(Or(str, list), Use(list_to_tuple)),
        Optional('label_shape', default=[1]): And(Or(str, list), Use(list_to_tuple)),
        Optional('low'): Or(
            float,
            And(int, Use(input_int_to_float)),
            And(list, Use(input_int_to_float)),
            And(str, Use(input_int_to_float))),
        Optional('high'): Or(
            float,
            And(int, Use(input_int_to_float)),
            And(list, Use(input_int_to_float)),
            And(str, Use(input_int_to_float))),
        Optional('dtype'): And(Or(str, list), Use(input_to_list)),
    },
    Optional('sparse_dummy_v2'): {
        'dense_shape': And(Or(str, list), Use(list_to_tuple)),
        Optional('label_shape', default=[1]): And(Or(str, list), Use(list_to_tuple)),
        Optional('sparse_ratio'): Or(
            float,
            And(list, Use(input_int_to_float)),
            And(int, Use(input_int_to_float))),
        Optional('low'): Or(
            float,
            And(int, Use(input_int_to_float)),
            And(list, Use(input_int_to_float)),
            And(str, Use(input_int_to_float))),
        Optional('high'): Or(
            float,
            And(int, Use(input_int_to_float)),
            And(list, Use(input_int_to_float)),
            And(str, Use(input_int_to_float))),
        Optional('dtype'): And(Or(str, list), Use(input_to_list)),
    },
 
    Optional('dummy'): {
        'shape': And(Or(str, list), Use(list_to_tuple)),
        Optional('low'): Or(
            float,
            And(int, Use(input_int_to_float)),
            And(list, Use(input_int_to_float)),
            And(str, Use(input_int_to_float))),
        Optional('high'): Or(
            float,
            And(int, Use(input_int_to_float)),
            And(list, Use(input_int_to_float)),
            And(str, Use(input_int_to_float))),
        Optional('dtype'): And(Or(str, list), Use(input_to_list)),
        Optional('label'): bool,
    },
    Optional('bert'): {
        'root': str,
        'label_file': str,
        Optional('task'): And(str, lambda s: s in ["classifier", "squad"]),
        Optional('model_type'): And(str, lambda s: s in ['bert', 'xlnet', 'xlm']),
    },
    Optional('mzbert'): {
        'root': str,
        'label_file': str,
        Optional('task'): And(str, lambda s: s in ["classifier", "squad"]),
        Optional('model_type'): And(str, lambda s: s in ['bert', 'xlnet', 'xlm']),
    },
    Optional('VOCRecord'): {
        'root': str,
    },
    Optional('COCORecord'): {
        'root': str,
        Optional('num_cores'): int,
    },
    Optional('COCORaw'): {
        'root': str,
        Optional('img_dir'): str,
        Optional('anno_dir'): str,
        Optional('num_cores'): int,
    },
    Optional('COCONpy'): {
        'root': str,
        'npy_dir': str,
        Optional('anno_dir'): str,
        Optional('num_cores'): int,
    },
    Optional('ImagenetRaw'): {
        'data_path': str,
        Optional('image_list'): str,
    },
    Optional('style_transfer'): {
        'content_folder': str,
        'style_folder': str,
        Optional('crop_ratio'): float,
        Optional('resize_shape'): And(Or(str, list), Use(input_to_list_int)),
        Optional('image_format'): str,
    },
    Optional('GLUE'): {
        'data_dir': str,
        'model_name_or_path': str,
        Optional('max_seq_length'): int,
        Optional('do_lower_case'): bool,
        Optional('task'): str,
        Optional('model_type'): str,
        Optional('dynamic_length'): bool,
        Optional('evaluate'): bool
    },
    # TO BE DEPRECATED!
    Optional('Imagenet'): {
        'root': str,
    },
})

dataloader_schema = Schema({
    Optional('last_batch', default='rollover'): And(str, lambda s: s in ['rollover', 'discard']),
    Optional('batch_size', default=None): And(int, lambda s: s > 0),
    'dataset': dataset_schema,
    Optional('filter'): filter_schema,
    Optional('transform'): transform_schema,
    Optional('shuffle', default = False): And(bool, lambda s: s in [True, False]),
    Optional('distributed', default = False): And(bool, lambda s: s in [True, False]),
})

configs_schema = Schema({
    Optional('cores_per_instance'): And(int, lambda s: s > 0),
    Optional('num_of_instance', default=1): And(int, lambda s: s > 0),
    Optional('inter_num_of_threads'): And(int, lambda s: s > 0),
    Optional('intra_num_of_threads'): And(int, lambda s: s > 0),
    Optional('kmp_blocktime'): And(int, lambda s: s >= 0),
    Optional('kmp_affinity', default='granularity=fine,verbose,compact,1,0'): str,
})

optimizer_schema = Schema({
    Optional('SGD'): {
        'learning_rate': Use(float),
        Optional('momentum'): Use(float),
        Optional('nesterov'): bool,
        Optional('weight_decay'): Use(float)
    },
    Optional('AdamW'): {
        'weight_decay': Use(float),
        Optional('learning_rate', default=0.001): Use(float),
        Optional('beta_1', default=0.9): Use(float),
        Optional('beta_2', default=0.999): Use(float),
        Optional('epsilon', default=1e-07): Use(float),
        Optional('amsgrad', default=False): bool
    },
    Optional('Adam'): {
        Optional('learning_rate', default=0.001): Use(float),
        Optional('beta_1', default=0.9): Use(float),
        Optional('beta_2', default=0.999): Use(float),
        Optional('epsilon', default=1e-07): Use(float),
        Optional('amsgrad', default=False): bool
    }, 
})

criterion_schema = Schema({
    Optional('CrossEntropyLoss'): {
        Optional('reduction', default='mean'): \
            And(str, lambda s: s in ['none', 'sum', 'mean', 'auto', 'sum_over_batch_size']),
        Optional('from_logits', default=False):
            And(bool, lambda s: s in [True, False]),
    },
    Optional('SparseCategoricalCrossentropy'): {
        Optional('reduction', default='mean'): \
            And(str, lambda s: s in ['none', 'sum', 'mean', 'auto', 'sum_over_batch_size']),
        Optional('from_logits', default=False):
            And(bool, lambda s: s in [True, False]),
    },
    Optional('KnowledgeDistillationLoss'): {
        Optional('temperature'): And(float, lambda s: s > 0),
        Optional('loss_types'): And(list, lambda s: all(i in ['CE', 'KL', 'MSE'] for i in s)),
        Optional('loss_weights'): And(list, lambda s: all(i >= 0 for i in s) and sum(s) == 1.0),
    },
    Optional('IntermediateLayersKnowledgeDistillationLoss'): {
        'layer_mappings':
            And(Or(tuple, list), lambda s: all(len(i) in [2, 4] for i in s)),
        Optional('loss_types'):
            And(Or(tuple, list), lambda s: all(i in ['MSE', 'KL', 'L1'] for i in s)),
        Optional('loss_weights'):
            And(Or(tuple, list), lambda s: all(i >= 0 for i in s)),
        Optional('add_origin_loss'): bool,
    }
})

train_schema = Schema({
    'optimizer': optimizer_schema,
    'criterion': criterion_schema,
    Optional('dataloader'): dataloader_schema,
    Optional('epoch', default=1): int,
    Optional('start_epoch', default=0): int,
    Optional('end_epoch'): int,
    Optional('iteration'): int,
    Optional('frequency'): int,
    Optional('execution_mode', default='eager'): And(str, lambda s: s in ['eager', 'graph']),
    Optional('postprocess'): {
        Optional('transform'): postprocess_schema
    },
    # TODO reserve for multinode training support
    Optional('hostfile'): str
})

weight_compression_schema = Schema({
    Optional('initial_sparsity', default=0): And(float, lambda s: s < 1.0 and s >= 0.0),
    Optional('target_sparsity', default=0.97): float,
    Optional('max_sparsity_ratio_per_layer', default=0.98):float,
    Optional('prune_type', default="basic_magnitude"): str,
    Optional('start_epoch', default=0): int,
    Optional('end_epoch', default=4): int,
    Optional('start_step', default=0): int,
    Optional('end_step', default=0): int,
    Optional('update_frequency', default=1.0): float,
    Optional('update_frequency_on_step', default=1):int,
    Optional('excluded_names', default=[]):list,
    Optional('prune_domain', default="global"): str,
    Optional('names', default=[]): list,
    Optional('extra_excluded_names', default=None): list,
    Optional('prune_layer_type', default=None): list,
    Optional('sparsity_decay_type', default="exp"): str,
    Optional('pattern', default="tile_pattern_1x1"): str,

    Optional('pruners'): And(list, \
                               lambda s: all(isinstance(i, Pruner) for i in s))
})

# weight_compression_pytorch_schema = Schema({},ignore_extra_keys=True)

approach_schema = Schema({
    Hook('weight_compression', handler=_valid_prune_sparsity): object,
    Hook('weight_compression_pytorch', handler=_valid_prune_sparsity): object,
    Optional('weight_compression'): weight_compression_schema,
    Optional('weight_compression_pytorch'): weight_compression_schema,
})

default_workspace = './nc_workspace/{}/'.format(
                                           datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))

COCOmAP_input_order_schema = Schema({
    Optional('num_detections'): int,
    'boxes': int,
    'scores': int,
    'classes': int
})

schema = Schema({
    'model': {
        'name': str,
        'framework': And(str, lambda s: s in list(FRAMEWORKS.keys()) + ['NA']),
        Optional('inputs', default=[]): And(Or(str, list), Use(input_to_list)),
        Optional('outputs', default=[]): And(Or(str, list), Use(input_to_list)),
    },
    Optional('version', default=float(__version__.split('.')[0])): And(
                                          Or(float,
                                             And(int, Use(input_int_to_float)),
                                             And(str, Use(input_int_to_float))),
                                          lambda s: s == float(__version__.split('.')[0])),
    Optional('device', default='cpu'): And(str, lambda s: s in ['cpu', 'gpu']),
    Optional('quantization', default={'approach': 'post_training_static_quant', \
                                      'calibration': {'sampling_size': [100]}, \
                                      'recipes': {'scale_propagation_max_pooling': True,
                                                      'scale_propagation_concat': True,
                                                      'first_conv_or_matmul_quantization': True},
                                      'model_wise': {'weight': {'bit': [7.0]},
                                                     'activation': {}},
                                      }): {
        Optional('approach', default='post_training_static_quant'): And(
            str,
            # TODO check if framework support dynamic quantize
            # Now only onnruntime and pytorch supoort
            lambda s: s in ['post_training_static_quant',
                            'post_training_dynamic_quant',
                            'post_training_auto_quant',
                            'quant_aware_training']),
        Optional('train', default=None): train_schema,
        Optional('advance', default=None): {
            Optional('bias_correction'): And(str, lambda s: s in ['weight_empirical']),
        },
        Optional('calibration', default={'sampling_size': [100]}): {
            Optional('sampling_size', default=[100]): And(Or(str, int, list), Use(input_to_list)),
            Optional('dataloader', default=None): dataloader_schema
        },
        Optional('recipes', default={'scale_propagation_max_pooling': True,
                                         'scale_propagation_concat': True,
                                         'first_conv_or_matmul_quantization': True}): {
            Optional('scale_propagation_max_pooling', default=True):
                    And(bool, lambda s: s in [True, False]),
            Optional('scale_propagation_concat', default=True):
                    And(bool, lambda s: s in [True, False]),
            Optional('first_conv_or_matmul_quantization', default=True):
                    And(bool, lambda s: s in [True, False]),
            Optional('fast_bias_correction', default=False):
                    And(bool, lambda s: s in [True, False]),
            Optional('weight_correction', default=False):
                    And(bool, lambda s: s in [True, False]),
        },
        Optional('model_wise', default={'weight': {'bit': [7.0]}, 'activation': {}}): {
            Optional('weight', default= {'bit': [7.0]}): {
                Optional('granularity', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['per_channel', 'per_tensor'] for i in s)),
                # asym_float are only for PyTorch framework
                Optional('scheme', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['asym', 'sym', 'asym_float'] for i in s)),
                Optional('dtype', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['int8', 'uint8', 'fp32', 'bf16', 'fp16'] for i in s)),
                Optional('algorithm', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['minmax'] for i in s)),
                Optional('bit', default=[7.0]):  And(
                    Or(float, list),
                    Use(input_to_list_float),
                    lambda s: all(0.0 < i <= 7.0 for i in s))

            },
            Optional('activation', default=None): {
                Optional('granularity', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['per_channel', 'per_tensor'] for i in s)),
                Optional('scheme', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['asym', 'sym'] for i in s)),
                Optional('dtype', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['int8', 'uint8', 'fp32', 'bf16', 'fp16'] for i in s)),
                # compute_dtypeis only for PyTorch framework
                Optional('compute_dtype', default=['uint8']): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['int8', 'uint8', 'fp32', 'bf16', 'None'] for i in s)),
                # placeholder are only for PyTorch framework
                Optional('algorithm', default=None): And(
                    Or(str, list),
                    Use(input_to_list),
                    lambda s: all(i in ['minmax', 'kl', 'placeholder'] for i in s)),
            }
        },
        Optional('optype_wise', default=None): {
            str: ops_schema
        },
        Optional('op_wise', default=None): {
            str: ops_schema
        },
    },

    Optional('graph_optimization'): graph_optimization_schema,
    Optional('mixed_precision'): mixed_precision_schema,

    Optional('model_conversion'): model_conversion_schema,

    Optional('tuning', default={
        'strategy': {'name': 'basic'},
        'accuracy_criterion': {'relative': 0.01, 'higher_is_better': True},
        'objective': 'performance',
        'exit_policy': {'timeout': 0, 'max_trials': 100, 'performance_only': False},
        'random_seed': 1978, 'tensorboard': False,
        'workspace': {'path': default_workspace},
        'diagnosis': False,
        }): {
        Optional('strategy', default={'name': 'basic'}): {
            'name': And(str, lambda s: s in STRATEGIES), Optional('sigopt_api_token'): str,
            Optional('sigopt_project_id'): str,
            Optional('sigopt_experiment_name', default='nc-tune'): str,
            Optional('accuracy_weight', default=1.0): float,
            Optional('latency_weight', default=1.0): float
        } ,
        Hook('accuracy_criterion', handler=_valid_accuracy_field): object,
        Optional('accuracy_criterion', default={'relative': 0.01}): {
            Optional('relative'): And(Or(str, float), Use(percent_to_float)),
            Optional('absolute'): And(Or(str, int, float), Use(percent_to_float)),
            Optional('higher_is_better', default=True): bool,
        },
        Optional('objective', default='performance'): And(str, lambda s: s in OBJECTIVES),
        Hook('multi_objectives', handler=_valid_multi_objectives): object,
        Optional('multi_objectives'):{ 
            Optional('objective'): And(
                Or(str, list), Use(input_to_list), lambda s: all(i in OBJECTIVES for i in s)),
            Optional('weight'): And(Or(str, list), Use(input_to_list_float)),
            Optional('higher_is_better'): And(
                Or(str, bool, list), Use(input_to_list_bool)),
        },
        Optional('exit_policy', default={'timeout': 0,
                                         'max_trials': 100,
                                         'performance_only': False}): {
            Optional('timeout', default=0): int,
            Optional('max_trials', default=100): int,
            Optional('performance_only', default=False): bool,
        },
        Optional('random_seed', default=1978): int,
        Optional('tensorboard', default=False): And(bool, lambda s: s in [True, False]),
        Optional('workspace', default={'path': default_workspace}): {
            Optional('path', default=None): str,
            Optional('resume'): str
        },
        Optional('diagnosis', default = {
            'diagnosis_after_tuning': False,
            'op_list': [],
            'iteration_list': [1],
            'inspect_type': 'activation',
            'save_to_disk': True,
            'save_path': './nc_workspace/inspect_saved/',
        }):{
            Optional('diagnosis_after_tuning', default=False): And(bool, lambda s: s in [True, False]),
            Optional('op_list', default=[]): And(Or(str, list), Use(input_to_list)),
            Optional('iteration_list', default=[1]): And(Or(int, list), Use(input_to_list_int)),
            Optional('inspect_type', default='all'): And(str, lambda s : s in ['all', 'activation', 'weight']),
            Optional('save_to_disk', default=True): And(bool, lambda s: s in [True, False]),
            Optional('save_path', default='./nc_workspace/inspect_saved/'): str,
        },
    },
    Optional('evaluation'): {
        Hook('accuracy', handler=_valid_multi_metrics): object,
        Optional('accuracy'): {
            Hook('multi_metrics', handler=_valid_metric_length): object,
            Optional('multi_metrics', default=None): {
                Optional('weight'): And(Or(str, list), Use(input_to_list_float)),
                Optional('higher_is_better'): And(
                    Or(str, list), Use(input_to_list_bool)),
                Optional('topk'): And(int, lambda s: s in [1, 5]),
                Optional('mAP'): {
                    Optional('anno_path'): str,
                    Optional('iou_thrs', default=0.5):
                            Or(And(str, lambda s: s in ['0.5:0.05:0.95']),
                               And(float, lambda s: s <= 1.0 and s >= 0.0)),
                    Optional('map_points', default=0): And(int, lambda s: s in [0, 11, 101])
                },
                Optional('COCOmAP'): {
                    Optional('anno_path'): str,
                    Optional('map_key', default='DetectionBoxes_Precision/mAP'): str
                },
                Optional('COCOmAPv2'): {
                    Optional('anno_path'): str,
                    Optional('map_key', default='DetectionBoxes_Precision/mAP'): str,
                    Optional('output_index_mapping', default={'num_detections': -1, 
                                                      'boxes': 0, 
                                                      'scores': 1, 
                                                      'classes': 2}): COCOmAP_input_order_schema
                },
                Optional('VOCmAP'): {
                    Optional('anno_path'): str
                },
                Optional('SquadF1'): Or({}, None),
                Optional('MSE'): {
                    Optional('compare_label'): bool
                },
                Optional('RMSE'): {
                    Optional('compare_label'): bool
                },
                Optional('MAE'): {
                    Optional('compare_label'): bool
                },
                Optional('Accuracy'): Or({}, None),
                Optional('Loss'): Or({}, None),
                Optional('BLEU'): Or({}, None),
                Optional('SquadF1'): Or({}, None),
                Optional('F1'): Or({}, None),
                Optional('mIOU'): {
                    Optional('num_classes'): int
                },
                Optional('GLUE'): {
                    Optional('task'): str
                },
                Optional('ROC'): {
                    Optional('task'): str
                },
            }, 
            Optional('metric', default=None): {
                Optional('topk'): And(int, lambda s: s in [1, 5]),
                Optional('mAP'): {
                    Optional('anno_path'): str,
                    Optional('iou_thrs', default=0.5):
                            Or(And(str, lambda s: s in ['0.5:0.05:0.95']),
                               And(float, lambda s: s <= 1.0 and s >= 0.0)),
                    Optional('map_points', default=0): And(int, lambda s: s in [0, 11, 101])
                },
                Optional('COCOmAP'): {
                    Optional('anno_path'): str,
                    Optional('map_key', default='DetectionBoxes_Precision/mAP'): str
                },
                Optional('COCOmAPv2'): {
                    Optional('anno_path'): str,
                    Optional('map_key', default='DetectionBoxes_Precision/mAP'): str,
                    Optional('output_index_mapping', default={'num_detections': -1, 
                                                      'boxes': 0, 
                                                      'scores': 1, 
                                                      'classes': 2}): COCOmAP_input_order_schema
                },
                Optional('VOCmAP'): {
                    Optional('anno_path'): str
                },
                Optional('SquadF1'): Or({}, None),
                Optional('MSE'): {
                    Optional('compare_label'): bool
                },
                Optional('RMSE'): {
                    Optional('compare_label'): bool
                },
                Optional('MAE'): {
                    Optional('compare_label'): bool
                },
                Optional('Accuracy'): Or({}, None),
                Optional('Loss'): Or({}, None),
                Optional('BLEU'): Or({}, None),
                Optional('SquadF1'): Or({}, None),
                Optional('F1'): Or({}, None),
                Optional('mIOU'): {
                    Optional('num_classes'): int
                },
                Optional('GLUE'): {
                    Optional('task'): str
                },
                Optional('ROC'): {
                    Optional('task'): str
                },
            },
            Optional('configs'): configs_schema,
            Optional('iteration', default=-1): int,
            Optional('dataloader'): dataloader_schema,
            Optional('postprocess'): {
                Optional('transform'): postprocess_schema
            },
        },
        Optional('performance'): {
            Optional('warmup', default=5): int,
            Optional('iteration', default=-1): int,
            Optional('configs'): configs_schema,
            Optional('dataloader'): dataloader_schema,
            Optional('postprocess'): {
                Optional('transform'): postprocess_schema
            }
        },
    },
    Optional('pruning'): {
        Hook('train', handler=_valid_prune_epoch): object,
        Optional("train"): train_schema,
        Optional("approach"): approach_schema
    },

    Optional('distillation'): {
        Optional("train"): train_schema
    },

    Optional('auto_distillation'): {
        "search": {
            "search_space": dict,
            Optional("search_algorithm", default=""): str,
            Optional("metrics", default=[]): list,
            Optional("higher_is_better", default=[]): list,
            Optional("max_trials", default=1): int,
            Optional("seed", default=42): int,
            },
        Optional("flash_distillation"): {
            Optional("knowledge_transfer"): {
                Optional("block_names", default=[]): list,
                "layer_mappings_for_knowledge_transfer": list,
                Optional("loss_types", default=[]): list,
                Optional("loss_weights", default=[]): list,
                Optional("add_origin_loss", default=[]): list,
                Optional("train_steps", default=[]): list,
                },
            Optional("regular_distillation"): {
                Optional("block_names", default=[]): list,
                "layer_mappings_for_knowledge_transfer": list,
                Optional("loss_types", default=[]): list,
                Optional("loss_weights", default=[]): list,
                Optional("add_origin_loss", default=[]): list,
                Optional("train_steps", default=[]): list,
                },
            },
    },

    Optional('nas'): {
        Optional("approach", default=None): Or(str, None),
        Optional("search"): {
            Optional("search_space", default=None): Or(dict, None),
            Optional("search_algorithm", default=None): Or(str, None),
            Optional("metrics", default=None): list,
            Optional("higher_is_better", default=None): list,
            Optional("max_trials", default=None): int,
            Optional("seed", default=42): int,
            },
        Optional("dynas"): {
            Optional("supernet", default=None): str,
            Optional("metrics", default=None): list,
            Optional("population", default=50): int,
            Optional("num_evals", default=100000): int,
            Optional("results_csv_path", default=None): str,
            Optional("dataset_path", default=None): str,
            Optional("batch_size", default=64): int,
            },
    },

    Optional("train"): train_schema
})

quantization_default_schema = Schema({
    Optional('model', default={'name': 'default_model_name', \
                               'framework': 'NA', \
                                'inputs': [], 'outputs': []}): dict,

    Optional('version', default=float(__version__.split('.')[0])): str,

    Optional('device', default='cpu'): str,

    Optional('quantization', default={'approach': 'post_training_static_quant', \
                                      'calibration': {'sampling_size': [100]},
                                      'recipes': {'scale_propagation_max_pooling': True,
                                                      'scale_propagation_concat': True,
                                                      'first_conv_or_matmul_quantization': True},
                                      'model_wise': {'weight': {'bit': [7.0]},
                                                     'activation': {}},
                                    }): dict,

    Optional('tuning', default={
        'strategy': {'name': 'basic'},
        'accuracy_criterion': {'relative': 0.01, 'higher_is_better': True},
        'objective': 'performance',
        'exit_policy': {'timeout': 0, 'max_trials': 100, 'performance_only': False},
        'random_seed': 1978, 'tensorboard': False,
        'workspace': {'path': default_workspace}}): dict,

    Optional('evaluation', default={'accuracy': {'metric': {'topk': 1}}}): dict
})

pruning_default_schema = Schema({
    Optional('model', default={'name': 'default_model_name', \
                               'framework': 'NA', \
                                'inputs': [], 'outputs': []}): dict,

    Optional('version', default=float(__version__.split('.')[0])): str,

    Optional('device', default='cpu'): str,

    Optional('tuning', default={
        'random_seed': 1978, 'tensorboard': False,
        'workspace': {'path': default_workspace}}): dict,

    Optional('pruning', default={'approach': {'weight_compression':{'initial_sparsity': 0.0, \
                                            'target_sparsity': 0.97, 'start_epoch': 0, \
                                            'end_epoch': 4}}}): dict,

    Optional('evaluation', default={'accuracy': {'metric': {'topk': 1}}}): dict
})

graph_optimization_default_schema = Schema({
    Optional('model', default={'name': 'resnet50', \
                               'framework': 'NA', \
                                'inputs': [], 'outputs': []}): dict,

    Optional('version', default=float(__version__.split('.')[0])): str,

    Optional('device', default='cpu'): str,

    Optional('quantization', default={'approach': 'post_training_static_quant', 
                                    'calibration': {'sampling_size': [100]},
                                    'recipes': {'scale_propagation_max_pooling': True,
                                                    'scale_propagation_concat': True,
                                                    'first_conv_or_matmul_quantization': True},
                                    'model_wise': {'weight': {'bit': [7.0]},
                                                    'activation': {}}}): dict,

    Optional('tuning', default={
        'strategy': {'name': 'basic'},
        'accuracy_criterion': {'relative': 0.01, 'higher_is_better': True},
        'objective': 'performance',
        'exit_policy': {'timeout': 0, 'max_trials': 100, 'performance_only': False},
        'random_seed': 1978, 'tensorboard': False,
        'workspace': {'path': default_workspace}}): dict,

    Optional('evaluation', default={'accuracy': {'metric': {'topk': 1}}}): dict,

    Optional('graph_optimization', default={'precisions': ['bf16, fp32']}): dict 
})

mixed_precision_default_schema = Schema({
    Optional('model', default={'name': 'resnet50', \
                               'framework': 'NA', \
                                'inputs': [], 'outputs': []}): dict,

    Optional('version', default=float(__version__.split('.')[0])): str,

    Optional('device', default='cpu'): str,

    Optional('quantization', default={'approach': 'post_training_static_quant', 
                                    'calibration': {'sampling_size': [100]},
                                    'recipes': {'scale_propagation_max_pooling': True,
                                                    'scale_propagation_concat': True,
                                                    'first_conv_or_matmul_quantization': True},
                                    'model_wise': {'weight': {'bit': [7.0]},
                                                    'activation': {}}}): dict,

    Optional('tuning', default={
        'strategy': {'name': 'basic'},
        'accuracy_criterion': {'relative': 0.01, 'higher_is_better': True},
        'objective': 'performance',
        'exit_policy': {'timeout': 0, 'max_trials': 100, 'performance_only': False},
        'random_seed': 1978, 'tensorboard': False,
        'workspace': {'path': default_workspace}}): dict,

    Optional('evaluation', default={'accuracy': {'metric': {'topk': 1}}}): dict,

    Optional('mixed_precision', default={'precisions': ['bf16, fp32']}): dict 
})

benchmark_default_schema = Schema({
    Optional('model', default={'name': 'resnet50', \
                               'framework': 'NA', \
                                'inputs': [], 'outputs': []}): dict,

    Optional('version', default=float(__version__.split('.')[0])): str,

    Optional('device', default='cpu'): str,

    Optional('quantization', default={'approach': 'post_training_static_quant', 
                                    'calibration': {'sampling_size': [100]},
                                    'recipes': {'scale_propagation_max_pooling': True,
                                                    'scale_propagation_concat': True,
                                                    'first_conv_or_matmul_quantization': True},
                                    'model_wise': {'weight': {'bit': [7.0]},
                                                    'activation': {}}}): dict,

    Optional('tuning', default={
        'strategy': {'name': 'basic'},
        'accuracy_criterion': {'relative': 0.01, 'higher_is_better': True},
        'objective': 'performance',
        'exit_policy': {'timeout': 0, 'max_trials': 100, 'performance_only': False},
        'random_seed': 1978, 'tensorboard': False,
        'workspace': {'path': default_workspace}}): dict,

    Optional('evaluation', default={'accuracy': {'metric': {'topk': 1}}}): dict
})

distillation_default_schema = Schema({
    Optional('model', default={'name': 'default_model_name', \
                               'framework': 'NA', \
                               'inputs': [], 'outputs': []}): dict,

    Optional('version', default=float(__version__.split('.')[0])): str,

    Optional('device', default='cpu'): str,

    Optional('tuning', default={
        'random_seed': 1978, 'tensorboard': False,
        'workspace': {'path': default_workspace}}): dict,

    Optional('distillation', default={
        'train': {'start_epoch': 0, 'end_epoch': 10, 
                  'iteration': 1000, 'frequency': 1, 
                  'optimizer': {'SGD': {'learning_rate': 0.001}}, 
                  'criterion': {'KnowledgeDistillationLoss': 
                                 {'temperature': 1.0, 
                                  'loss_types': ['CE', 'KL'], 
                                  'loss_weights': [0.5, 0.5]}}}}): dict,

    Optional('evaluation', default={'accuracy': {'metric': {'topk': 1}}}):dict
 
})

class Conf(object):
    """config parser.

    Args:
        cfg_fname (string): The path to the configuration file.

    """
    def __init__(self, cfg_fname):
        assert cfg_fname is not None
        self.usr_cfg = DotDict(self._read_cfg(cfg_fname))

    def _read_cfg(self, cfg_fname):
        """Load a config file following yaml syntax.

           Args:
               cfg_fname(string): The name of configuration yaml file
        """
        try:
            with open(cfg_fname, 'r') as f:
                content = f.read()
                cfg = yaml.safe_load(content)
                validated_cfg = schema.validate(cfg)

            # if user yaml doesn't include version field, neural_compressor will write a supported version
            # into it.
            if 'version' not in cfg:
                leading_whitespace = re.search(r"[ \t]*model\s*:",
                                               content).group().split("model")[0]
                content = re.sub(r'model\s*:',
                                 'version: {}\n\n{}model:'.format(
                                                               float(__version__.split('.')[0]),
                                                               leading_whitespace
                                                           ),
                                 content)
                with open(cfg_fname, 'w') as f:
                    f.write(content)

            return validated_cfg   
        except FileNotFoundError as f:
            logger.error("{}.".format(f))
            raise RuntimeError(
                "The yaml file is not exist. Please check the file name or path."
            )
        except Exception as e:
            logger.error("{}.".format(e))
            raise RuntimeError(
                "The yaml file format is not correct. Please refer to document."
            )

    def map_pyconfig_to_cfg(self, pythonic_config):
        mapping = {
            'device': pythonic_config.quantization.device,
            'model.inputs': pythonic_config.quantization.inputs,
            'model.outputs': pythonic_config.quantization.outputs,
            'model.framework': pythonic_config.quantization.backend,
            'quantization.approach': pythonic_config.quantization.approach,
            'quantization.calibration.sampling_size': 
                pythonic_config.quantization.calibration_sampling_size,
            'quantization.optype_wise': pythonic_config.quantization.op_type_list,
            'quantization.op_wise': pythonic_config.quantization.op_type_list,
            'distillation.train.criterion': pythonic_config.distillation.criterion,
            'distillation.train.optimizer': pythonic_config.distillation.optimizer,
            'pruning.approach.weight_compression': pythonic_config.pruning.weight_compression,
            'nas.approach': pythonic_config.nas.approach,
            'nas.search': pythonic_config.nas.search,
            'nas.dynas': pythonic_config.nas.dynas,
            'tuning.strategy.name': pythonic_config.quantization.strategy,
            'tuning.accuracy_criterion.relative': 
                pythonic_config.quantization.accuracy_criterion.relative,
            'tuning.accuracy_criterion.absolute':
                pythonic_config.quantization.accuracy_criterion.absolute,
            'tuning.accuracy_criterion.higher_is_better':
                pythonic_config.quantization.accuracy_criterion.higher_is_better,
            'tuning.objective': pythonic_config.quantization.objective,
            'tuning.exit_policy.timeout': pythonic_config.quantization.timeout,
            'tuning.exit_policy.max_trials': pythonic_config.quantization.max_trials,
            'tuning.exit_policy.performance_only': pythonic_config.quantization.performance_only,
            'tuning.random_seed': pythonic_config.options.random_seed,
            'tuning.workspace.path': pythonic_config.options.workspace,
            'tuning.workspace.resume': pythonic_config.options.resume_from,
            'evaluation.performance.warmup': pythonic_config.benchmark.warmup,
            'evaluation.performance.iteration': pythonic_config.benchmark.iteration,
            'evaluation.performance.configs.cores_per_instance':
                pythonic_config.benchmark.cores_per_instance,
            'evaluation.performance.configs.num_of_instance':
                pythonic_config.benchmark.num_of_instance,
            'evaluation.performance.configs.inter_num_of_threads':
                pythonic_config.benchmark.inter_num_of_threads,
            'evaluation.performance.configs.intra_num_of_threads':
                pythonic_config.benchmark.intra_num_of_threads,
            'evaluation.accuracy.configs.cores_per_instance':
                pythonic_config.benchmark.cores_per_instance,
            'evaluation.accuracy.configs.num_of_instance':
                pythonic_config.benchmark.num_of_instance,
            'evaluation.accuracy.configs.inter_num_of_threads':
                pythonic_config.benchmark.inter_num_of_threads,
            'evaluation.accuracy.configs.intra_num_of_threads':
                pythonic_config.benchmark.intra_num_of_threads,
            'use_bf16': pythonic_config.quantization.use_bf16,
            'reduce_range': pythonic_config.quantization.reduce_range
        }

        for k, v in mapping.items():
            if k in ['tuning.accuracy_criterion.relative', 'tuning.accuracy_criterion.absolute']:
                target_key = str(pythonic_config.quantization.accuracy_criterion)
                if target_key not in k and 'accuracy_criterion' in self.usr_cfg.tuning:
                    if target_key in self.usr_cfg.tuning.accuracy_criterion and \
                                    k.split('.')[-1] in self.usr_cfg.tuning.accuracy_criterion:
                        self.usr_cfg.tuning.accuracy_criterion.pop(k.split('.')[-1])
                    continue
            if v is not None:
                deep_set(self.usr_cfg, k, v)

    def _convert_cfg(self, src, dst):
        """Helper function to merge user defined dict into default dict.

           If the key in src doesn't exist in dst, then add this key and value
           pair to dst.
           If the key in src is in dst, then override the value in dst with the
           value in src.

        Args:
            src (dict): The source dict merged from
            dst (dict): The source dict merged to

        Returns:
            dict: The merged dict from src to dst
        """
        for key in src:
            if key in dst:
                if isinstance(dst[key], dict) and isinstance(src[key], dict):
                    if key in ['accuracy_criterion', 'metric', 'dataset', 
                        'criterion', 'optimizer']:
                        # accuracy_criterion can only have one of absolute and relative
                        # others can only have one item
                        inter_key = src[key].keys() & dst[key].keys()-{'higher_is_better'}
                        if len(inter_key) == 0:
                            dst[key] = {}
                    if key == 'accuracy' and src[key].get('multi_metrics', None):
                        dst[key].pop('metric', None)
                    self._convert_cfg(src[key], dst[key])
                elif dst[key] == src[key]:
                    pass  # same leaf value
                else:
                    dst[key] = src[key]
            elif isinstance(src[key], dict):
                dst[key] = DotDict(self._convert_cfg(src[key], {}))
            else:
                dst[key] = src[key]
        return dst

class Quantization_Conf(Conf):
    """config parser.

    Args:
        cfg: The path to the configuration file or DotDict object or None.

    """

    def __init__(self, cfg=None):
        if isinstance(cfg, str):
            self.usr_cfg = DotDict(self._read_cfg(cfg))
        elif isinstance(cfg, DotDict):
            self.usr_cfg = DotDict(schema.validate(self._convert_cfg(
                cfg, copy.deepcopy(quantization_default_schema.validate(dict())))))
        else:
            self.usr_cfg = DotDict(quantization_default_schema.validate(dict()))
        self._model_wise_tune_space = None
        self._opwise_tune_space = None

    def _merge_dicts(self, src, dst):
        """Helper function to merge src dict into dst dict.

           If the key in src doesn't exist in dst, then skip
           If the key in src is in dst and the value intersects with the one in
           dst, then override the value in dst with the intersect value.

        Args:
            src (dict): The source dict merged from
            dst (dict): The source dict merged to

        Returns:
            dict: The merged dict from src to dst
        """
        for key in src:
            if key in dst:
                if isinstance(dst[key], dict) and isinstance(src[key], dict):
                    self._merge_dicts(src[key], dst[key])
                elif dst[key] == src[key] or src[key] is None:
                    pass  # same leaf value
                else:
                    value = [value for value in src[key]
                             if value in dst[key] or isinstance(value, float)]
                    if value != []:
                        dst[key] = value

        return dst

    def modelwise_tune_space(self, model_wise_quant):
        cfg = self.usr_cfg

        self._model_wise_tune_space = OrderedDict()
        for optype in model_wise_quant.keys():
            if cfg.quantization.optype_wise and optype in cfg.quantization.optype_wise:
                self._model_wise_tune_space[optype] = self._merge_dicts(
                    cfg.quantization.optype_wise[optype],
                    model_wise_quant[optype])
            else:
                self._model_wise_tune_space[optype] = self._merge_dicts(
                    cfg.quantization.model_wise,
                    model_wise_quant[optype])

        return self._model_wise_tune_space

class Pruning_Conf(Conf):
    """config parser.

    Args:
        cfg: The path to the configuration file or DotDict object or None.

    """

    def __init__(self, cfg=None):
        if isinstance(cfg, str):
            self._read_cfg(cfg)
            self.usr_cfg = DotDict(self._read_cfg(cfg))
        elif isinstance(cfg, DotDict):
            self.usr_cfg = DotDict(schema.validate(self._convert_cfg(
                cfg, copy.deepcopy(pruning_default_schema.validate(dict())))))
        else:
            self.usr_cfg = DotDict(pruning_default_schema.validate(dict()))

class Graph_Optimization_Conf(Quantization_Conf):
    """config parser.

    Args:
        cfg: The path to the configuration file or DotDict object or None.

    """

    def __init__(self, cfg=None):
        if isinstance(cfg, str):
            self.usr_cfg = DotDict(self._read_cfg(cfg))
        elif isinstance(cfg, DotDict):
            self.usr_cfg = DotDict(schema.validate(self._convert_cfg(
                cfg, copy.deepcopy(graph_optimization_default_schema.validate(dict())))))
        else:
            self.usr_cfg = DotDict(graph_optimization_default_schema.validate(dict()))

class MixedPrecision_Conf(Quantization_Conf):
    """config parser.

    Args:
        cfg: The path to the configuration file or DotDict object or None.

    """

    def __init__(self, cfg=None):
        if isinstance(cfg, str):
            self.usr_cfg = DotDict(self._read_cfg(cfg))
        elif isinstance(cfg, DotDict):
            self.usr_cfg = DotDict(self._convert_cfg(
                cfg, copy.deepcopy(mixed_precision_default_schema.validate(dict()))))
        else:
            self.usr_cfg = DotDict(mixed_precision_default_schema.validate(dict()))

class Benchmark_Conf(Conf):
    """config parser.

    Args:
        cfg: The path to the configuration file or DotDict object or None.

    """

    def __init__(self, cfg=None):
        if isinstance(cfg, str):
            self.usr_cfg = DotDict(self._read_cfg(cfg))
        elif isinstance(cfg, DotDict):
            self.usr_cfg = DotDict(schema.validate(self._convert_cfg(
                cfg, copy.deepcopy(benchmark_default_schema.validate(dict())))))
        else:
            self.usr_cfg = DotDict(benchmark_default_schema.validate(dict()))

class Distillation_Conf(Conf):
    """config parser.

    Args:
        cfg: The path to the configuration file or DotDict object or None.

    """

    def __init__(self, cfg=None):
        if isinstance(cfg, str):
            self.usr_cfg = DotDict(self._read_cfg(cfg))
        elif isinstance(cfg, DotDict):
            self.usr_cfg = DotDict(schema.validate(self._convert_cfg(
                cfg, copy.deepcopy(distillation_default_schema.validate(dict())))))
        else:
            self.usr_cfg = DotDict(distillation_default_schema.validate(dict()))

class NASConfig(Conf):
    """config parser.

    Args:
        approach: The approach of the NAS.
        search_algorithm: The search algorithm for NAS procedure.

    """

    def __init__(self, approach=None, search_space=None, search_algorithm=None):
        assert approach is None or (isinstance(approach, str) and approach), \
            "Should set approach with a string."
        usr_cfg = {
            'model': {'name': 'nas', 'framework': 'NA'},
            'nas': {'approach': approach,
                    'search': {
                        'search_space': search_space,
                        'search_algorithm': search_algorithm}
                    },
        }
        self.usr_cfg = DotDict(usr_cfg)
        if approach and approach != 'basic':
            self.usr_cfg.nas[approach] = DotDict({})
            self.__setattr__(approach, self.usr_cfg.nas[approach])

    def validate(self):
        self.usr_cfg = schema.validate(self.usr_cfg)
        
    @property
    def nas(self):
        return self.usr_cfg.nas

class DefaultConf(DotDict):
    def __getitem__(self, key):
        if key not in self:
            self[key] = DefaultConf({})
        value = self.get(key, None)
        return value

    __getattr__ = __getitem__

conf = DefaultConf({})
QuantConf = Quantization_Conf
PruningConf = Pruning_Conf
GraphOptConf = Graph_Optimization_Conf
BenchmarkConf = Benchmark_Conf
DistillationConf = Distillation_Conf
MixedPrecisionConf = MixedPrecision_Conf
