#  Copyright 2022 The HuggingFace Team. 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 logging
import os
from functools import reduce
from pathlib import Path
from typing import Any, Optional, Union

from transformers.utils import TRANSFORMERS_CACHE, is_offline_mode

import yaml
from huggingface_hub import hf_hub_download
from neural_compressor.conf.config import Conf, Distillation_Conf, Pruning_Conf, Quantization_Conf
from optimum.intel.neural_compressor.utils import CONFIG_NAME


logger = logging.getLogger(__name__)


class IncConfig:
    CONFIG_NAME = "best_configure.yml"

    def __init__(self, config_path: str):
        """
        Args:
            config_path (`str`):
                Path to the YAML configuration file used to control the tuning behavior.
        Returns:
            config: IncConfig object.
        """

        self.path = config_path
        self.config = Conf(config_path)
        self.usr_cfg = self.config.usr_cfg

    def get_config(self, keys: str):
        return reduce(lambda d, key: d.get(key) if d else None, keys.split("."), self.usr_cfg)

    def set_config(self, keys: str, value: Any):
        d = self.usr_cfg
        keys = keys.split(".")
        for key in keys[:-1]:
            d = d.setdefault(key, {})
        d[keys[-1]] = value

    def set_tolerance(self, tolerance_criterion: Union[int, float]):
        if not isinstance(tolerance_criterion, (int, float)):
            raise TypeError(
                f"Supported type for performance tolerance are int and float, got {type(tolerance_criterion)}"
            )
        if "absolute" in self.get_config("tuning.accuracy_criterion"):
            self.set_config("tuning.accuracy_criterion.absolute", tolerance_criterion)
        else:
            if not -1 < tolerance_criterion < 1:
                raise ValueError("Relative performance tolerance must not be <=-1 or >=1.")
            self.set_config("tuning.accuracy_criterion.relative", tolerance_criterion)

    @classmethod
    def from_pretrained(cls, config_name_or_path: str, config_file_name: Optional[str] = None, **kwargs):
        """
        Instantiate an IncConfig object from a configuration file which can either be hosted on
        huggingface.co or from a local directory path.

        Args:
            config_name_or_path (`str`):
                Repository name in the Hugging Face Hub or path to a local directory containing the configuration file.
            config_file_name (`str`, *optional*):
                Name of the configuration file.
            cache_dir (`str`, *optional*):
                Path to a directory in which a downloaded configuration should be cached if the standard cache should
                not be used.
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force to (re-)download the configuration files and override the cached versions if
                they exist.
            resume_download (`bool`, *optional*, defaults to `False`):
                Whether or not to delete incompletely received file. Attempts to resume the download if such a file
                exists.
            revision(`str`, *optional*):
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                git-based system for storing models and other artifacts on huggingface.co, so ``revision`` can be any
                identifier allowed by git.
        Returns:
            config: IncConfig object.
        """

        cache_dir = kwargs.get("cache_dir", None)
        force_download = kwargs.get("force_download", False)
        resume_download = kwargs.get("resume_download", False)
        revision = kwargs.get("revision", None)

        config_file_name = config_file_name if config_file_name is not None else cls.CONFIG_NAME

        if os.path.isdir(config_name_or_path):
            config_file = os.path.join(config_name_or_path, config_file_name)
        elif os.path.isfile(config_name_or_path):
            config_file = config_name_or_path
        else:
            local_files_only = False
            if is_offline_mode():
                logger.info("Offline mode: forcing local_files_only=True")
                local_files_only = True
            if cache_dir is None:
                cache_dir = TRANSFORMERS_CACHE
            if isinstance(cache_dir, Path):
                cache_dir = str(cache_dir)
            try:
                config_file = hf_hub_download(
                    repo_id=config_name_or_path,
                    filename=config_file_name,
                    revision=revision,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
                )
            except EnvironmentError as err:
                logger.error(err)
                msg = (
                    f"Can't load config for '{config_name_or_path}'. Make sure that:\n\n"
                    f"-'{config_name_or_path}' is a correct model identifier listed on 'https://huggingface.co/models'\n\n"
                    f"-or '{config_name_or_path}' is a correct path to a directory containing a {config_file_name} file\n\n"
                )

                if revision is not None:
                    msg += (
                        f"- or '{revision}' is a valid git identifier (branch name, a tag name, or a commit id) that "
                        f"exists for this model name as listed on its model page on 'https://huggingface.co/models'\n\n"
                    )

                raise EnvironmentError(msg)
        config = cls(config_file)

        return config


class IncOptimizedConfig(IncConfig):
    def __init__(self, config_path: str):
        """
        Args:
            config_path (`str`):
                Path to the YAML configuration file used to control the tuning behavior.
        """

        self.path = config_path
        self.config = self._read_config()
        self.usr_cfg = self.config

    def _read_config(self):
        with open(self.path, "r") as f:
            try:
                config = yaml.load(f, Loader=yaml.Loader)
            except yaml.YAMLError as err:
                logger.error(err)

        return config


class IncQuantizationConfig(IncConfig):
    CONFIG_NAME = "quantization.yml"

    def __init__(self, config_path: str):
        """
        The configuration describing the quantization process.

        Args:
            config_path (`str`):
                Path to the YAML configuration file used to control the tuning behavior.
        """

        self.path = config_path
        self.config = Quantization_Conf(config_path)
        self.usr_cfg = self.config.usr_cfg


class IncPruningConfig(IncConfig):
    CONFIG_NAME = "prune.yml"

    def __init__(self, config_path: str):
        """
        The configuration describing the pruning process.

        Args:
            config_path (`str`):
                Path to the YAML configuration file used to control the tuning behavior.
        """

        self.path = config_path
        self.config = Pruning_Conf(config_path)
        self.usr_cfg = self.config.usr_cfg


class IncDistillationConfig(IncConfig):
    CONFIG_NAME = "distillation.yml"

    def __init__(self, config_path: str):
        """
        Args:
            config_path (:obj:`str`):
                Path to the YAML configuration file used to control the tuning behavior.
        """

        self.path = config_path
        self.config = Distillation_Conf(config_path)
        self.usr_cfg = self.config.usr_cfg
