# 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 paddle
import json
import warnings
from paddle.fluid import core

__all__ = ['set_config']


def set_config(config=None):
    r"""
    Set the configuration for kernel, layout and dataloader auto-tuning.

    1. kernel: When it is enable, exhaustive search method will be used to select
    and cache the best algorithm for the operator in the tuning iteration. Tuning
    parameters are as follows:

    - enable(bool): Whether to enable kernel tuning.
    - tuning_range(list): Start and end iteration for auto-tuning. Default: [1, 10].

    2. layout: When it is enable, the best data layout such as NCHW or NHWC will be
    determined based on the device and data type. When the origin layout setting is
    not best, layout transformation will be automaticly performed to improve model
    performance. Layout auto-tuning only supports dygraph mode currently. Tuning
    parameters are as follows:

    - enable(bool): Whether to enable layout tuning.

    3. dataloader: When it is enable, the best num_workers will be selected to replace
    the origin dataloader setting. Tuning parameters are as follows:

    - enable(bool): Whether to enable dataloader tuning.

    Args:
        config (dict|str|None, optional): Configuration for auto-tuning. If it is a
            dictionary, the key is the tuning type, and the value is a dictionary
            of the corresponding tuning parameters. If it is a string, the path of
            a json file will be specified and the tuning configuration will be set
            by the the json file. Default: None, auto-tuning for kernel, layout and
            dataloader will be enabled.

    Examples:
        .. code-block:: python
            :name: auto-tuning

            import paddle
            import json

            # config is a dict.
            config = {
                "kernel": {
                    "enable": True,
                    "tuning_range": [1, 5],
                },
                "layout": {
                    "enable": True,
                },
                "dataloader": {
                    "enable": True,
                }
            }
            paddle.incubate.autotune.set_config(config)

            # config is the path of json file.
            config_json = json.dumps(config)
            with open('config.json', 'w') as json_file:
                json_file.write(config_json)
            paddle.incubate.autotune.set_config('config.json')

    """
    if config is None:
        core.enable_autotune()
        core.enable_layout_autotune()
        paddle.fluid.reader.set_autotune_config(use_autotune=True)
        return

    config_dict = {}
    if isinstance(config, dict):
        config_dict = config
    elif isinstance(config, str):
        try:
            with open(config, 'r') as filehandle:
                config_dict = json.load(filehandle)
        except Exception as e:
            print('Load config error: {}'.format(e))
            warnings.warn("Use default configuration for auto-tuning.")

    if "kernel" in config_dict:
        kernel_config = config_dict["kernel"]
        if "enable" in kernel_config:
            if isinstance(kernel_config['enable'], bool):
                if kernel_config['enable']:
                    core.enable_autotune()
                else:
                    core.disable_autotune()
            else:
                warnings.warn(
                    "The auto-tuning configuration of the kernel is incorrect."
                    "The `enable` should be bool. Use default parameter instead."
                )
        if "tuning_range" in kernel_config:
            if isinstance(kernel_config['tuning_range'], list):
                tuning_range = kernel_config['tuning_range']
                assert len(tuning_range) == 2
                core.set_autotune_range(tuning_range[0], tuning_range[1])
            else:
                warnings.warn(
                    "The auto-tuning configuration of the kernel is incorrect."
                    "The `tuning_range` should be list. Use default parameter instead."
                )
    if "layout" in config_dict:
        layout_config = config_dict["layout"]
        if "enable" in layout_config:
            if isinstance(layout_config['enable'], bool):
                if layout_config['enable']:
                    core.enable_layout_autotune()
                else:
                    core.disable_layout_autotune()
            else:
                warnings.warn(
                    "The auto-tuning configuration of the layout is incorrect."
                    "The `enable` should be bool. Use default parameter instead."
                )
    if "dataloader" in config_dict:
        dataloader_config = config_dict["dataloader"]
        use_autoune = False
        if "enable" in dataloader_config:
            if isinstance(dataloader_config['enable'], bool):
                use_autoune = dataloader_config['enable']
            else:
                warnings.warn(
                    "The auto-tuning configuration of the dataloader is incorrect."
                    "The `enable` should be bool. Use default parameter instead."
                )
        if "tuning_steps" in dataloader_config:
            if isinstance(dataloader_config['tuning_steps'], int):
                paddle.fluid.reader.set_autotune_config(
                    use_autoune, dataloader_config['tuning_steps'])
            else:
                warnings.warn(
                    "The auto-tuning configuration of the dataloader is incorrect."
                    "The `tuning_steps` should be int. Use default parameter instead."
                )
                paddle.fluid.reader.set_autotune_config(use_autoune)