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

from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject
from deepspeed.profiling.constants import *


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

        self.enabled = None
        self.profile_step = None
        self.module_depth = None
        self.top_modules = None

        if FLOPS_PROFILER in param_dict.keys():
            flops_profiler_dict = param_dict[FLOPS_PROFILER]
        else:
            flops_profiler_dict = {}

        self._initialize(flops_profiler_dict)

    def _initialize(self, flops_profiler_dict):
        self.enabled = get_scalar_param(flops_profiler_dict,
                                        FLOPS_PROFILER_ENABLED,
                                        FLOPS_PROFILER_ENABLED_DEFAULT)

        self.profile_step = get_scalar_param(flops_profiler_dict,
                                             FLOPS_PROFILER_PROFILE_STEP,
                                             FLOPS_PROFILER_PROFILE_STEP_DEFAULT)

        self.module_depth = get_scalar_param(flops_profiler_dict,
                                             FLOPS_PROFILER_MODULE_DEPTH,
                                             FLOPS_PROFILER_MODULE_DEPTH_DEFAULT)

        self.top_modules = get_scalar_param(flops_profiler_dict,
                                            FLOPS_PROFILER_TOP_MODULES,
                                            FLOPS_PROFILER_TOP_MODULES_DEFAULT)

        self.detailed = get_scalar_param(flops_profiler_dict,
                                         FLOPS_PROFILER_DETAILED,
                                         FLOPS_PROFILER_DETAILED_DEFAULT)

        self.output_file = get_scalar_param(flops_profiler_dict,
                                            FLOPS_PROFILER_OUTPUT_FILE,
                                            FLOPS_PROFILER_OUTPUT_FILE_DEFAULT)