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

"""Provide a layer of abstraction for an OpenVINO runtime environment."""

import logging
from typing import Dict, List, Union

import numpy as np
from openvino import Core, IENetwork, Blob, DataPtr

from openvino.exceptions import UserInputError
from openvino.impl import Function, Node, PartialShape, Type
from openvino.opset1.ops import result
from openvino.utils.types import NumericData, get_shape, get_dtype

import tests

log = logging.getLogger(__name__)


def runtime(backend_name: str = "CPU") -> "Runtime":
    """Create a Runtime object (helper factory)."""
    return Runtime(backend_name)


def get_runtime():
    """Return runtime object."""
    if tests.BACKEND_NAME is not None:
        return runtime(backend_name=tests.BACKEND_NAME)
    else:
        return runtime()


def _convert_inputs(cnn_network: IENetwork) -> None:
    """WA converts unsupported input images formats."""
    precision_map = {
        "FP64": "FP32",
        "I64": "I32",
        "U32": "I32",
    }

    for cnn_input in cnn_network.input_info:
        try:
            _precision = precision_map[cnn_network.input_info[cnn_input].precision]
            cnn_network.input_info[cnn_input].precision = _precision
        except KeyError:
            pass


def _convert_val(val):
    """WA converts unsupported input values."""
    if type(val) is np.ndarray:
        if val.dtype == np.float64:
            return np.array(val, dtype=np.float32)
        elif val.dtype == np.int64:
            return np.array(val, dtype=np.int32)
        return np.array(val)

    return np.array(val, dtype=np.float32)


def apply_ng_type(output: DataPtr, ng_type: Type):
    ng_ie_supported_type_map = {
        Type.boolean.get_type_name(): "BOOL",
        Type.f32.get_type_name(): "FP32",
        Type.i8.get_type_name(): "I8",
        Type.i32.get_type_name(): "I32",
        Type.u8.get_type_name(): "U8",
    }
    if ng_type.get_type_name() in ng_ie_supported_type_map:
        output.precision = ng_ie_supported_type_map[ng_type.get_type_name()]


class Runtime(object):
    """Represents an nGraph runtime environment."""

    def __init__(self, backend_name: str) -> None:
        self.backend_name = backend_name
        log.debug("Creating Inference Engine for %s" % backend_name)
        self.backend = Core()
        assert backend_name in self.backend.available_devices, (
            'The requested device "' + backend_name + '" is not supported!'
        )

    def set_config(self, config: Dict[str, str]) -> None:
        """Set the inference engine configuration."""
        self.backend.set_config(config, device_name=self.backend_name)

    def __repr__(self) -> str:
        return "<Runtime: Backend='{}'>".format(self.backend_name)

    def computation(self, node_or_function: Union[Node, Function], *inputs: Node) -> "Computation":
        """Return a callable Computation object."""
        if isinstance(node_or_function, Node):
            ng_function = Function(node_or_function, inputs, node_or_function.name)
            return Computation(self, ng_function)
        elif isinstance(node_or_function, Function):
            return Computation(self, node_or_function)
        else:
            raise TypeError(
                "Runtime.computation must be called with an nGraph Function object "
                "or an nGraph node object an optionally Parameter node objects. "
                "Called with: %s",
                node_or_function,
            )


class Computation(object):
    """nGraph callable computation object."""

    def __init__(self, runtime: Runtime, ng_function: Function) -> None:
        self.runtime = runtime
        self.function = ng_function
        self.parameters = ng_function.get_parameters()
        self.results = ng_function.get_results()
        self.network_cache = {}

    def __repr__(self) -> str:
        params_string = ", ".join([param.name for param in self.parameters])
        return "<Computation: {}({})>".format(self.function.get_name(), params_string)

    def _get_ie_output_blob_name(self, outputs: Dict, ng_result: result) -> str:
        if len(self.results) == 1:
            return next(iter(outputs.keys()))
        else:
            prev_layer = ng_result.input(0).get_source_output()
            out_name = prev_layer.get_node().get_friendly_name()
            if prev_layer.get_node().get_output_size() != 1:
                out_name += "." + str(prev_layer.get_index())
            return out_name

    def _get_ie_output_blob_buffer(self, output_blobs: Dict[str, Blob], ng_result: result) -> np.ndarray:
        out_name = self._get_ie_output_blob_name(output_blobs, ng_result)
        out_blob = output_blobs[out_name]

        if out_blob.tensor_desc.layout == "SCALAR":
            return out_blob.buffer.reshape(())
        else:
            return out_blob.buffer

    def convert_buffers(self, source_buffers, target_dtypes):
        converted_buffers = []
        for i in range(len(source_buffers)):
            target_dtype = target_dtypes[i]
            # custom conversion for bf16
            if self.results[i].get_output_element_type(0) == Type.bf16:
                converted_buffers.append((source_buffers[i].view(np.uint32) >> 16).astype(np.uint16))
            else:
                converted_buffers.append(source_buffers[i].astype(target_dtype))
        return converted_buffers

    def __call__(self, *input_values: NumericData) -> List[NumericData]:
        """Run computation on input values and return result."""
        # Input validation
        if len(input_values) < len(self.parameters):
            raise UserInputError(
                "Expected %s params, received not enough %s values.", len(self.parameters), len(input_values)
            )
        # ignore not needed input values
        input_values = input_values[:len(self.parameters)]

        input_values = [_convert_val(input_value) for input_value in input_values]
        input_shapes = [get_shape(input_value) for input_value in input_values]

        param_names = [param.friendly_name for param in self.parameters]

        if self.network_cache.get(str(input_shapes)) is None:
            cnn_network = IENetwork(self.function)
            if self.function.is_dynamic():
                cnn_network.reshape(dict(zip(param_names, input_shapes)))
            # Convert unsupported inputs of the network
            _convert_inputs(cnn_network)
            self.network_cache[str(input_shapes)] = cnn_network
        else:
            cnn_network = self.network_cache[str(input_shapes)]

        # set output blobs precission based on nG results
        for ng_result in self.results:
            ie_out_name = self._get_ie_output_blob_name(cnn_network.outputs, ng_result)
            apply_ng_type(cnn_network.outputs[ie_out_name], ng_result.get_output_element_type(0))

        executable_network = self.runtime.backend.load_network(cnn_network, self.runtime.backend_name)

        for parameter, input in zip(self.parameters, input_values):
            parameter_shape = parameter.get_output_partial_shape(0)
            input_shape = PartialShape(input.shape)
            if len(input.shape) > 0 and not parameter_shape.compatible(input_shape):
                raise UserInputError(
                    "Provided tensor's shape: %s does not match the expected: %s.",
                    input_shape,
                    parameter_shape,
                )

        request = executable_network.create_infer_request()
        request.infer(dict(zip(param_names, input_values)))

        # Set order of output blobs compatible with nG Function
        result_buffers = [self._get_ie_output_blob_buffer(request.output_blobs, result)
                          for result in self.results]

        # Since OV overwrite result data type we have to convert results to the original one.
        original_dtypes = [get_dtype(result.get_output_element_type(0)) for result in self.results]
        converted_buffers = self.convert_buffers(result_buffers, original_dtypes)
        return converted_buffers
