text: #
  # Copyright 2016 The BigDL Authors.
  #
  # 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.
  #
  
  from bigdl.dllib.nn.layer import Identity, Model
  from .ops_mapping import _convert_map as convert_map
  from .converter_utils import parse_node_attr, parse_tensor_data
  from bigdl.dllib.utils.log4Error import *
  
  
  class OnnxLoader(object):
  
      def load_model(self, file_path):
          import onnx
          model_proto = onnx.load_model(file_path)
          # self._ir_version = model_proto.ir_version
          # self._opset_import = model_proto.opset_import
          # self._producer_name = model_proto.producer_name
          # self._producer_version = model_proto.producer_version
          # self._domain = model_proto.domain
          # self._model_version = model_proto.model_version
          # self._doc_string = model_proto.doc_string
          graph_proto = model_proto.graph
          return self.load_graph(graph_proto)
  
      def load_graph(self, graph_proto):
          if not graph_proto:
              invalidInputError(False, "Graph proto is required")
  
          input_nodes = list()
          output_nodes = list()
          tensor_map = dict()
          initialized_tensors = set()
          module_map = dict()
          root_nodes = list()
          dummy_root = Identity()()
  
          for tensor in graph_proto.initializer:
              if not tensor.name.strip():
                  invalidInputError(False, "Tensor's name is required")
              initialized_tensors.add(tensor.name)
              tensor_data = parse_tensor_data(tensor)
              tensor_map[tensor.name] = (tensor_data, tensor_data.shape)
  
          for gin in graph_proto.input:
              if gin.name not in initialized_tensors:
                  input_nodes.append(gin.name)
                  shape = tuple([dim.dim_value for dim in gin.type.tensor_type.shape.dim])
                  module_map[gin.name] = Identity()(dummy_root)
                  tensor_map[gin.name] = (None, shape)
  
          for gout in graph_proto.output:
              if gout.name not in initialized_tensors:
                  output_nodes.append(gout.name)
  
          for node in graph_proto.node:
              name = node.name.strip()
              op_type = node.op_type
              inputs = [tensor_map[n] for n in node.input]
              outputs = node.output
              prev_modules = [module_map[n] for n in node.input if n not in initialized_tensors]
              attrs = parse_node_attr(node)
  
              if len(prev_modules) == 0:
                  root_nodes.append((name, op_type))
                  prev_modules = [dummy_root]
  
              bigdl_module, outputs_shape = self._make_module_from_onnx_node(op_type, inputs,
                                                                             prev_modules, attrs,
                                                                             outputs)
  
              invalidInputError(len(outputs) == len(outputs_shape),
                                f"size of outputs {len(outputs)} doesn't match outputs_shape"
                                f" ${len(outputs_shape)}")
  
              for out, out_shape in zip(outputs, outputs_shape):
                  module_map[out] = bigdl_module
                  tensor_map[out] = (None, out_shape)
  
          in_modules = [module_map[m] for m in input_nodes]
          out_modules = [module_map[m] for m in output_nodes]
          model = Model([dummy_root], out_modules)
  
          return model
  
      def _make_module_from_onnx_node(self, op_type, inputs, prev_modules, attrs, outputs):
          module = None
          out_shapes = []
          if op_type in convert_map:
              module, out_shapes = convert_map[op_type](inputs, prev_modules, attrs, outputs)
          else:
              invalidInputError(False, "not implement yet")
          return module, out_shapes
  
  
  def load(model_path):
      loader = OnnxLoader()
      return loader.load_model(model_path)
  
  
  def load_model_proto(model_proto):
      loader = OnnxLoader()
      return loader.load_graph(model_proto.graph)
  