#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# 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 copy
import os

from ...strategy import Strategy

_tuning_supported_passes = ["sharding", "recompute"]


def _get_pass_config(strategy, pass_name):
    config = getattr(strategy, pass_name)
    return config


class TuningConfig:
    """
    A uniform config wrap:
    distributed strategy: the user defined configuration for optimization pass
    tuning config: configuration for the tuning process: mode (profile or cost model), log dir, extra tuning config for optimization like search range for specific
    """

    def __init__(self, strategy):

        if not isinstance(strategy, Strategy):
            raise TypeError("'strategy' must be object of class `Strategy`.")

        self._tuning_passes_name = set()
        self._dist_strategy = copy.deepcopy(strategy)
        self._mode = None
        self._profile_start_step = None
        self._profile_end_step = None
        self._project_dir = None
        self._max_num_trial = None
        self._early_stop = None
        self._debug = None

        self._initialize()

    @property
    def mode(self):
        return self._mode

    @property
    def profile_start_step(self):
        return self._profile_start_step

    @property
    def profile_end_step(self):
        return self._profile_end_step

    @property
    def project_dir(self):
        return self._project_dir

    @property
    def tuning_passes_name(self):
        return self._tuning_passes_name

    @property
    def max_num_trial(self):
        return self._max_num_trial

    @property
    def early_stop(self):
        return self._early_stop

    @property
    def debug(self):
        return self._debug

    @property
    def dist_strategy(self):
        return self._dist_strategy

    # initialize config with user define value or default value
    def _initialize(self):
        tuning_strategy = self._dist_strategy.tuning

        self._mode = tuning_strategy.get("mode", "PROFILE")
        self._profile_start_step = tuning_strategy.get("profile_start_step", 10)
        self._profile_end_step = tuning_strategy.get("profile_end_step", 30)
        self._max_num_trial = tuning_strategy.get("max_num_trial", 50)
        self._early_stop = tuning_strategy.get("early_stop", None)
        self._debug = tuning_strategy.get("debug", False)

        project_dir = tuning_strategy.get("project_dir", None)
        if not project_dir:
            project_dir = os.path.join(os.getcwd(), "OptimizationTuning")
        self._project_dir = project_dir

        for p in _tuning_supported_passes:
            if (
                getattr(self._dist_strategy, p)
                and _get_pass_config(self._dist_strategy, p).enable_tuning
            ):
                # TODO distinguish different args of each passes
                self._tuning_passes_name.add(p)

                p_strategy = getattr(self._dist_strategy, p)
                self.__dict__[p] = p_strategy

                # # TODO verify the user defined configs
                # tuning_config_for_pass = tuning_strategy.get(p, None)
                # if tuning_config_for_pass:
                #     for k, v in tuning_config_for_pass.items():
                #         self.__dict__[p][k] = v

    # (NOTE)tuning config ONLY wraps dist strategy for pass config which is to be tuned
    def __getattr__(self, item):
        return getattr(self._dist_strategy, item)