'''Copyright The Microsoft DeepSpeed Team'''
"""
Copyright (c) Microsoft Corporation
Licensed under the MIT license.
"""

from deepspeed.runtime.config_utils import get_scalar_param, get_dict_param, DeepSpeedConfigObject
from deepspeed.autotuning.constants import *


class DeepSpeedAutotuningConfig(DeepSpeedConfigObject):
    def __init__(self, param_dict):
        super(DeepSpeedAutotuningConfig, self).__init__()

        self.enabled = None
        self.start_step = None
        self.end_step = None
        self.metric_path = None
        self.arg_mappings = None
        self.metric = None
        self.model_info = None
        self.results_dir = None
        self.exps_dir = None
        self.overwrite = None

        if param_dict and AUTOTUNING in param_dict.keys():
            autotuning_dict = param_dict[AUTOTUNING]
        else:
            autotuning_dict = {}

        self._initialize(autotuning_dict)

    def _initialize(self, autotuning_dict):
        self.enabled = get_scalar_param(autotuning_dict,
                                        AUTOTUNING_ENABLED,
                                        AUTOTUNING_ENABLED_DEFAULT)

        self.fast = get_scalar_param(autotuning_dict,
                                     AUTOTUNING_FAST,
                                     AUTOTUNING_FAST_DEFAULT)

        self.results_dir = get_scalar_param(autotuning_dict,
                                            AUTOTUNING_RESULTS_DIR,
                                            AUTOTUNING_RESULTS_DIR_DEFAULT)
        assert self.results_dir, "results_dir cannot be empty"
        self.exps_dir = get_scalar_param(autotuning_dict,
                                         AUTOTUNING_EXPS_DIR,
                                         AUTOTUNING_EXPS_DIR_DEFAULT)
        assert self.exps_dir, "exps_dir cannot be empty"
        self.overwrite = get_scalar_param(autotuning_dict,
                                          AUTOTUNING_OVERWRITE,
                                          AUTOTUNING_OVERWRITE_DEFAULT)

        self.start_profile_step = get_scalar_param(
            autotuning_dict,
            AUTOTUNING_START_PROFILE_STEP,
            AUTOTUNING_START_PROFILE_STEP_DEFAULT)

        self.end_profile_step = get_scalar_param(autotuning_dict,
                                                 AUTOTUNING_END_PROFILE_STEP,
                                                 AUTOTUNING_END_PROFILE_STEP_DEFAULT)

        self.metric = get_scalar_param(autotuning_dict,
                                       AUTOTUNING_METRIC,
                                       AUTOTUNING_METRIC_DEFAULT)

        self.metric_path = get_scalar_param(autotuning_dict,
                                            AUTOTUNING_METRIC_PATH,
                                            AUTOTUNING_METRIC_PATH_DEFAULT)

        self.tuner_type = get_scalar_param(autotuning_dict,
                                           AUTOTUNING_TUNER_TYPE,
                                           AUTOTUNING_TUNER_TYPE_DEFAULT)

        self.tuner_early_stopping = get_scalar_param(
            autotuning_dict,
            AUTOTUNING_TUNER_EARLY_STOPPING,
            AUTOTUNING_TUNER_EARLY_STOPPING_DEFAULT)

        self.tuner_num_trials = get_scalar_param(autotuning_dict,
                                                 AUTOTUNING_TUNER_NUM_TRIALS,
                                                 AUTOTUNING_TUNER_NUM_TRIALS_DEFAULT)

        self.arg_mappings = get_dict_param(autotuning_dict,
                                           AUTOTUNING_ARG_MAPPINGS,
                                           AUTOTUNING_ARG_MAPPINGS_DEFAULT)

        self.model_info = get_model_info_config(autotuning_dict)

        self.model_info_path = get_scalar_param(autotuning_dict,
                                                AUTOTUNING_MODEL_INFO_PATH,
                                                AUTOTUNING_MODEL_INFO_PATH_DEFAULT)
        self.mp_size = get_scalar_param(autotuning_dict,
                                        AUTOTUNING_MP_SIZE,
                                        AUTOTUNING_MP_SIZE_DEFAULT)

        self.max_train_batch_size = get_dict_param(
            autotuning_dict,
            AUTOTUNING_MAX_TRAIN_BATCH_SIZE,
            AUTOTUNING_MAX_TRAIN_BATCH_SIZE_DEFAULT)

        self.min_train_batch_size = get_dict_param(
            autotuning_dict,
            AUTOTUNING_MIN_TRAIN_BATCH_SIZE,
            AUTOTUNING_MIN_TRAIN_BATCH_SIZE_DEFAULT)

        self.max_train_micro_batch_size_per_gpu = get_dict_param(
            autotuning_dict,
            AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU,
            AUTOTUNING_MAX_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT)

        self.min_train_micro_batch_size_per_gpu = get_dict_param(
            autotuning_dict,
            AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU,
            AUTOTUNING_MIN_TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT)

        self.num_tuning_micro_batch_sizes = get_dict_param(
            autotuning_dict,
            AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES,
            AUTOTUNING_NUM_TUNING_MICRO_BATCH_SIZES_DEFAULT)


def get_model_info_config(param_dict):
    if MODEL_INFO in param_dict and param_dict[MODEL_INFO] is not None:
        model_info_config = {}
        for key, default_value in MODEL_INFO_KEY_DEFAULT_DICT.items():
            model_info_config[key] = get_scalar_param(param_dict[MODEL_INFO],
                                                      key,
                                                      default_value)
        return model_info_config
    return None


def get_default_model_info_config():
    return MODEL_INFO_KEY_DEFAULT_DICT