# Copyright (C) 2018-2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import itertools
import os
import warnings
from pathlib import Path

import numpy as np
from common.constants import test_device, test_precision
from common.layer_utils import IEInfer
from mo.utils.ir_engine.ir_engine import IREngine

from common.utils.common_utils import generate_ir
from common.utils.parsers import mapping_parser


class CommonLayerTest:
    input_model_key = "input_model"

    def produce_model_path(self, framework_model, save_path):
        pass

    def get_framework_results(self, inputs_dict, model_path):
        pass

    def _test(self, framework_model, ref_net, ie_device, precision, ir_version, temp_dir, infer_timeout=60,
              enabled_transforms='', disabled_transforms='', **kwargs):
        """
        :param enabled_transforms/disabled_transforms: string with idxs of transforms that should be enabled/disabled.
                                                       Example: "transform_1,transform_2"
        """
        model_path = self.produce_model_path(framework_model=framework_model, save_path=temp_dir)

        # TODO Pass environment variables via subprocess environment
        os.environ['MO_ENABLED_TRANSFORMS'] = enabled_transforms
        os.environ['MO_DISABLED_TRANSFORMS'] = disabled_transforms

        mo_params = {self.input_model_key: model_path,
                     "output_dir": temp_dir,
                     "data_type": precision, "model_name": 'model'
                     }

        if 'input_shapes' in kwargs and len(kwargs['input_shapes']):
            input_shapes_str = []
            for ishape in kwargs['input_shapes']:
                input_shapes_str.append('[' + ','.join([str(i) for i in ishape]) + ']')
            mo_params.update(dict(input_shape=','.join(input_shapes_str)))

        if 'input_names' in kwargs and len(kwargs['input_names']):
            mo_params.update(dict(input=','.join(kwargs['input_names'])))

        exit_code, stderr = generate_ir(**mo_params)

        del os.environ['MO_ENABLED_TRANSFORMS']
        del os.environ['MO_DISABLED_TRANSFORMS']
        assert not exit_code, ("IR generation failed with {} exit code: {}".format(exit_code, stderr))

        path_to_xml = Path(temp_dir, 'model.xml')
        path_to_bin = Path(temp_dir, 'model.bin')

        # TODO: need to update ref graphs or get rid of this comparison
        # if ref_net is not None:
        #     ir = IREngine(path_to_xml, path_to_bin, precision=precision)
        #     (flag, resp) = ir.compare(ref_net)
        #     assert flag, '\n'.join(resp)

        from openvino.inference_engine import IECore
        core = IECore()
        net = core.read_network(path_to_xml, path_to_bin)
        inputs_info = {}
        for item in net.input_info.items():
            inputs_info[item[0]] = item[1].tensor_desc.dims

        # Prepare feed dict
        if 'kwargs_to_prepare_input' in kwargs and kwargs['kwargs_to_prepare_input']:
            inputs_dict = self._prepare_input(inputs_info, kwargs['kwargs_to_prepare_input'])
        else:
            inputs_dict = self._prepare_input(inputs_info)

        # IE infer:
        ie_engine = IEInfer(model=path_to_xml,
                            weights=path_to_bin,
                            device=ie_device)
        infer_res = ie_engine.infer(input_data=inputs_dict, infer_timeout=infer_timeout)

        if hasattr(self, 'skip_framework') and self.skip_framework:
            warnings.warn('Framework is skipped')
            return

        # Framework infer:
        fw_res = self.get_framework_results(inputs_dict=inputs_dict, model_path=model_path)

        if len(fw_res) == len(infer_res) == 1:
            # match output layers directly
            mapping_dict = {next(iter(fw_res)): next(iter(infer_res))}
        else:
            # Load mapping file
            mapping_dict = mapping_parser(path_to_xml.with_suffix('.mapping'))

        if 'custom_eps' in kwargs and kwargs['custom_eps'] is not None:
            custom_eps = kwargs['custom_eps']
        else:
            custom_eps = 1e-4

        # Compare Ie results with Framework results
        fw_eps = custom_eps if precision == 'FP32' else 5e-2
        assert self.compare_ie_results_with_framework(infer_res=infer_res, framework_res=fw_res,
                                                      mapping_dict=mapping_dict, framework_eps=fw_eps), \
            "Comparing with Framework failed: ie_res={}; framework_res={}.".format(infer_res, fw_res)

    # Feed dict for each input is filled with random number.
    # It is possible to redefine this function and generate your own input
    def _prepare_input(self, inputs_dict):
        for input in inputs_dict.keys():
            inputs_dict[input] = np.random.randint(-255, 255, inputs_dict[input]).astype(np.float32)
        return inputs_dict

    def compare_ie_results_with_framework(self, infer_res, framework_res, mapping_dict, framework_eps):
        is_ok = True
        from common.utils.common_utils import allclose
        for framework_out_name in framework_res:
            if framework_out_name not in mapping_dict:
                raise RuntimeError("Output {} not found in mapping file!".format(framework_out_name))

            ie_out_name = mapping_dict[framework_out_name]

            if not allclose(infer_res[ie_out_name], framework_res[framework_out_name], atol=framework_eps,
                            rtol=framework_eps):
                is_ok = False
                print("Max diff is {}".format(
                    np.array(abs(infer_res[ie_out_name] - framework_res[framework_out_name])).max()))
            else:
                print("Accuracy validation successful!\n")
                print("absolute eps: {}, relative eps: {}".format(framework_eps, framework_eps))
        return is_ok


def get_params(ie_device=None, precision=None):
    """
    :param ie_device: list of devices
    :param precision: list of precisions
    """

    ie_device_params = ie_device if ie_device else test_device
    precision_params = precision if precision else test_precision

    test_args = []
    for element in itertools.product(ie_device_params, precision_params):
        if element[0] == 'CPU' and element[1] == 'FP16':
            continue
        test_args.append(element)
    return test_args


def check_ir_version(left, right, ir_version):
    try:
        _ir_version = int(ir_version)
    except ValueError:
        raise RuntimeError("Wrong ir version type: {}, must be an integer".format(ir_version))
    left_bound = _ir_version - 1 if left is None else left
    right_bound = _ir_version + 1 if right is None else right
    return left_bound <= _ir_version < right_bound
