import logging
import os
from typing import Dict

import numpy as np

from .device_util import CPU_DEVICE, get_cpu_procs


class OpenvinoServable:
    """
    openvino servable 输入dict[str, ndarray]，输出dict[str, ndarray]
    """

    def __init__(self, local_path: str):
        self.local_path = local_path
        self.model_ins = None
        self.model_init_stat = False

        self.device = CPU_DEVICE

        self.model_input_list = None

    def load(self):
        if not self.model_init_stat:
            from openvino.inference_engine import IECore

            # --------------------------- 1. Read IR Generated by ModelOptimizer (.xml and .bin files) ------------
            ie_core = IECore()
            model_xml = self.local_path
            model_bin = os.path.splitext(model_xml)[0] + ".bin"
            net = ie_core.read_network(model=model_xml, weights=model_bin)
            # -----------------------------------------------------------------------------------------------------

            # ------------- 2. Load Plugin for inference engine and extensions library if specified --------------
            supported_layers = ie_core.query_network(net, self.device)
            not_supported_layers = [l for l in net.layers.keys() if l not in supported_layers]
            if len(not_supported_layers) != 0:
                logging.error("Following layers are not supported by the plugin for specified device %s: %s",
                              self.device, ', '.join(not_supported_layers))
                exit()
            # -----------------------------------------------------------------------------------------------------

            # ---------------------- 3. instantiate an executable version of the model ----------------------------
            available_procs = str(get_cpu_procs())
            self.model_ins = ie_core.load_network(network=net,
                                                  device_name=self.device,
                                                  config={"CPU_THREADS_NUM": available_procs})
            self.model_init_stat = True
            # -----------------------------------------------------------------------------------------------------

            # --------------------------- 4. Configure input & output ---------------------------------------------
            self.model_input_list = []
            for input_obj_name in self.model_ins.inputs:
                self.model_input_list.append(input_obj_name)
                logging.info("model input name: %s shape: %s",
                             input_obj_name, self.model_ins.inputs[input_obj_name].shape)
            # -----------------------------------------------------------------------------------------------------

    def forward(self, inputs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        if not isinstance(inputs, dict):
            raise Exception("openvino servable can only support dict input, got (%s) type" % str(type(inputs)))

        input_dict = {}
        for key in self.model_input_list:
            if key not in inputs:
                raise Exception("necessary input (%s) missing and model needs %s" % (key, self.model_input_list))
            input_dict[key] = inputs[key]

        output_dict = self.model_ins.infer(input_dict)

        return output_dict
