"""
Copyright (c) Microsoft Corporation
Licensed under the MIT license.
"""

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


class DeepSpeedFlopsProfilerConfig(object):
    def __init__(self, param_dict):
        """
        docstring
        """
        super(DeepSpeedFlopsProfilerConfig, self).__init__()

        self.enabled = None
        self.start_step = None
        self.end_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):
        """
        docstring
        """
        self.enabled = get_scalar_param(flops_profiler_dict,
                                        FLOPS_PROFILER_ENABLED,
                                        FLOPS_PROFILER_ENABLED_DEFAULT)

        self.start_step = get_scalar_param(flops_profiler_dict,
                                           FLOPS_PROFILER_START_STEP,
                                           FLOPS_PROFILER_START_STEP_DEFAULT)

        self.end_step = get_scalar_param(flops_profiler_dict,
                                         FLOPS_PROFILER_END_STEP,
                                         FLOPS_PROFILER_END_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)
