#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021 Intel Corporation
#
# 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 numpy as np
from collections import OrderedDict
from .strategy import strategy_registry, TuneStrategy
from ..utils import logger

@strategy_registry
class BasicTuneStrategy(TuneStrategy):
    """The basic tuning strategy which tunes the low precision model with below order.

    1. modelwise tuning for all quantizable ops.
    2. fallback tuning from bottom to top to decide the priority of which op has biggest impact
       on accuracy.
    3. incremental fallback tuning by fallbacking multiple ops with the order got from #2.

    Args:
        model (object):                        The FP32 model specified for low precision tuning.
        conf (Class):                          The Conf class instance initialized from user yaml
                                               config file.
        q_dataloader (generator):              Data loader for calibration, mandatory for
                                               post-training quantization.
                                               It is iterable and should yield a tuple (input,
                                               label) for calibration dataset containing label,
                                               or yield (input, _) for label-free calibration
                                               dataset. The input could be a object, list, tuple or
                                               dict, depending on user implementation, as well as
                                               it can be taken as model input.
        q_func (function, optional):           Reserved for future use.
        eval_dataloader (generator, optional): Data loader for evaluation. It is iterable
                                               and should yield a tuple of (input, label).
                                               The input could be a object, list, tuple or dict,
                                               depending on user implementation, as well as it can
                                               be taken as model input. The label should be able
                                               to take as input of supported metrics. If this
                                               parameter is not None, user needs to specify
                                               pre-defined evaluation metrics through configuration
                                               file and should set "eval_func" parameter as None.
                                               Tuner will combine model, eval_dataloader and
                                               pre-defined metrics to run evaluation process.
        eval_func (function, optional):        The evaluation function provided by user.
                                               This function takes model as parameter, and
                                               evaluation dataset and metrics should be
                                               encapsulated in this function implementation and
                                               outputs a higher-is-better accuracy scalar value.

                                               The pseudo code should be something like:

                                               def eval_func(model):
                                                    input, label = dataloader()
                                                    output = model(input)
                                                    accuracy = metric(output, label)
                                                    return accuracy
        dicts (dict, optional):                The dict containing resume information.
                                               Defaults to None.

    """

    def __init__(self, model, conf, q_dataloader, q_func=None,
                 eval_dataloader=None, eval_func=None, dicts=None, q_hooks=None):
        super(
            BasicTuneStrategy,
            self).__init__(
            model,
            conf,
            q_dataloader,
            q_func,
            eval_dataloader,
            eval_func,
            dicts,
            q_hooks)

    def next_tune_cfg(self):
        """The generator of yielding next tuning config to traverse by concrete strategies
           according to last tuning result.

        """
        # Model wise tuning
        op_cfgs = {}
        best_cfg = None
        if len(self.metric_name) == 1 or self.metric_weight is not None:
            best_acc = float('-inf') if self.higher_is_better else float('inf')
        else:
            best_acc = [float('-inf') if higher_is_better else float('inf') for \
                higher_is_better in self.metric_criterion]

        logger.debug("Start basic strategy by model-wise tuning.")
        for i, iterations in enumerate(self.calib_iter):
            op_cfgs['calib_iteration'] = int(iterations)
            op_cfgs['calib_sampling_size'] = int(self.calib_sampling_size[i])

            for combined_cfg in self.combined_model_wise_quant_cfgs:
                op_cfgs['op'] = OrderedDict()
                for op, op_cfg in self.opwise_quant_cfgs.items():
                    if op[1] in combined_cfg.keys() and len(op_cfg) > 0:
                        op_cfgs['op'][op] = copy.deepcopy(
                            self._get_common_cfg(combined_cfg[op[1]], op_cfg))
                    else:
                        op_cfgs['op'][op] = copy.deepcopy(
                            self.opwise_tune_cfgs[op][0])

                yield op_cfgs
                acc, _ = self.last_tune_result
                if not isinstance(acc, list) and ((self.higher_is_better and acc >= best_acc) \
                    or (not self.higher_is_better and acc <= best_acc)):
                    best_acc = acc
                    best_cfg = copy.deepcopy(op_cfgs)
                elif len(self.metric_name) > 1 and self.metric_weight is not None:
                    acc = np.mean(np.array(acc) * self.metric_weight)
                    if (self.higher_is_better and acc >= best_acc) or \
                        (not self.higher_is_better and acc <= best_acc):
                        best_acc = acc
                        best_cfg = copy.deepcopy(op_cfgs)
                elif len(self.metric_name) > 1 and self.metric_weight is None:
                    if all([acc_i >= best_i if higher_is_better else acc_i <= best_i for \
                        acc_i, best_i, higher_is_better in \
                        zip(acc, best_acc, self.metric_criterion)]):
                        best_acc = acc
                        best_cfg = copy.deepcopy(op_cfgs)

        if best_cfg is not None:
            fallback_dtypes = []
            for data_type in ["bf16", "fp16", "fp32"]:
                for op_type, tune_space in self.modelwise_tune_space.items():
                    if (data_type in tune_space["activation"]["dtype"] and
                        data_type not in fallback_dtypes):
                        fallback_dtypes.append(data_type)

            for fallback_dtype in fallback_dtypes:
                logger.debug(
                    "Continue basic strategy by sorting opwise {} fallback priority.".format
                    (fallback_dtype))
                ops_acc = OrderedDict()
                for op, configs in reversed(self.opwise_tune_cfgs.items()):
                    op_cfgs = copy.deepcopy(best_cfg)
                    for cfg in configs:
                        if ('weight' in cfg and fallback_dtype == cfg['weight']['dtype']) \
                           or ('weight' not in cfg
                               and fallback_dtype == cfg['activation']['dtype']):
                            if op_cfgs['op'][op]['weight'] is not None:
                                op_cfgs['op'][op]['weight'].clear()
                                op_cfgs['op'][op]['weight']['dtype'] = fallback_dtype
                            assert cfg['activation']['dtype'] == fallback_dtype
                            op_cfgs['op'][op]['activation'].clear()
                            op_cfgs['op'][op]['activation']['dtype'] = fallback_dtype
                    yield op_cfgs
                    acc, _ = self.last_tune_result
                    ops_acc[op] = acc

                op_cfgs = copy.deepcopy(best_cfg)
                if ops_acc is not None:
                    ordered_ops = sorted(ops_acc.keys(), key=lambda key: ops_acc[key],
                                         reverse=self.higher_is_better)
                    for op in ordered_ops:
                        old_cfg = copy.deepcopy(op_cfgs['op'][op])
                        for cfg in self.opwise_tune_cfgs[op]:
                            if ('weight' in cfg and fallback_dtype == cfg['weight']['dtype']) \
                               or ('weight' not in cfg
                                   and fallback_dtype == cfg['activation']['dtype']):
                                if op_cfgs['op'][op]['weight'] is not None:
                                    op_cfgs['op'][op]['weight'].clear()
                                    op_cfgs['op'][op]['weight']['dtype'] = fallback_dtype
                                assert cfg['activation']['dtype'] == fallback_dtype
                                op_cfgs['op'][op]['activation'].clear()
                                op_cfgs['op'][op]['activation']['dtype'] = fallback_dtype
                        yield op_cfgs
                        acc, _ = self.last_tune_result

                        if not isinstance(acc, list):
                            if ((self.higher_is_better and acc <= best_acc) \
                                or (not self.higher_is_better and acc >= best_acc)):
                                op_cfgs['op'][op] = copy.deepcopy(old_cfg)
                            else:
                                best_acc = acc
                        elif len(self.metric_name) > 1 and self.metric_weight is not None:
                            acc = np.mean(np.array(acc) * self.metric_weight)
                            if (self.higher_is_better and acc <= best_acc) or \
                                (not self.higher_is_better and acc >= best_acc):
                                op_cfgs['op'][op] = copy.deepcopy(old_cfg)
                            else:
                                best_acc = acc
                        elif len(self.metric_name) > 1 and self.metric_weight is None:
                            if all([acc_i >= best_i if higher_is_better else acc_i <= best_i for \
                                acc_i, best_i, higher_is_better in \
                                zip(acc, best_acc, self.metric_criterion)]):
                                op_cfgs['op'][op] = copy.deepcopy(old_cfg)
                            else:
                                best_acc = acc
 
                    op_cfgs = copy.deepcopy(best_cfg)
                    for op in ordered_ops:
                        for cfg in self.opwise_tune_cfgs[op]:
                            if ('weight' in cfg and fallback_dtype == cfg['weight']['dtype']) \
                               or ('weight' not in cfg
                                   and fallback_dtype == cfg['activation']['dtype']):
                                if op_cfgs['op'][op]['weight'] is not None:
                                    op_cfgs['op'][op]['weight'].clear()
                                    op_cfgs['op'][op]['weight']['dtype'] = fallback_dtype
                                assert cfg['activation']['dtype'] == fallback_dtype
                                op_cfgs['op'][op]['activation'].clear()
                                op_cfgs['op'][op]['activation']['dtype'] = fallback_dtype
                        yield op_cfgs
        else:
            logger.debug(self.opwise_tune_cfgs)
            op_cfgs['op'] = OrderedDict()
            for op in self.opwise_tune_cfgs.keys():
                op_cfgs['op'][op] = copy.deepcopy(self.opwise_tune_cfgs[op][0])
            yield op_cfgs

        return
