text: #  Copyright (c) 2020, Apple Inc. All rights reserved.
  #
  #  Use of this source code is governed by a BSD-3-clause license that can be
  #  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
  
  import os as _os
  import shutil as _shutil
  import stat as _stat
  import tempfile as _tempfile
  import warnings as _warnings
  
  from . import InputType, ImageType
  from .mil.passes.apply_common_pass_pipeline import apply_common_pass_pipeline
  from coremltools.converters._profile_utils import _profile
  from coremltools.converters.mil.mil import Builder as mb
  from coremltools.converters.mil.mil.passes.quantization_passes import AbstractQuantizationPass
  from coremltools.converters.mil.mil.types.symbolic import k_used_symbols, k_num_internal_syms
  from coremltools.models import MLModel
  from coremltools.models.model import _create_mlpackage
  from coremltools.models.utils import _MLMODEL_EXTENSION, _MLPACKAGE_EXTENSION
  
  try:
      from coremltools.libmodelpackage import ModelPackage
  except ModuleNotFoundError:
      pass
  
  
  class ConverterRegistry:
      frontends = {}
      backends = {}
      backend_alias_names = {}
  
      @staticmethod
      def frontend(converter):
          ConverterRegistry.frontends[converter.name] = converter
          return converter
  
      @staticmethod
      def backend(converter):
          ConverterRegistry.backends[converter.name] = converter
          if 'alias_names' in converter.__dict__:
              for name in converter.alias_names:
                  ConverterRegistry.backend_alias_names[name] = converter.name
          return converter
  
  
  @ConverterRegistry.frontend
  class MILFrontend:
      name = "milinternal"
  
      def __call__(self, model, *args, **kwargs):
          if "inputs" in kwargs and kwargs["inputs"] is not None:
              inputs = kwargs["inputs"]
              if not isinstance(inputs, (list, tuple)):
                  raise ValueError(
                      "Type of inputs should be list or tuple, got {} instead.".format(
                          type(inputs)
                      )
                  )
              if not all([isinstance(i, InputType) for i in inputs]):
                  raise ValueError(
                      "Type of inputs should be list or tuple of TensorType or ImageType, got {} instead.".format(
                          [type(i) for i in inputs]
                      )
                  )
  
              for idx, inp in enumerate(inputs):
                  # We set the default image format in MIL as NCHW, since only NCHW is
                  # natively supported by MIL ops (ex. Conv/Pool/etc.)
                  if isinstance(inp, ImageType) and inputs[idx].channel_first is None:
                      inputs[idx].channel_first = True
              model.set_main_input_types(tuple(inputs))
          return model
  
  
  @ConverterRegistry.frontend
  class TensorFlowFrontend:
      name = "tensorflow"
  
      def __call__(self, *args, **kwargs):
          from .frontend.tensorflow.load import TF1Loader
  
          tf1_loader = TF1Loader(*args, **kwargs)
          return tf1_loader.load()
  
  
  @ConverterRegistry.frontend
  class TensorFlow2Frontend:
      name = "tensorflow2"
  
      def __call__(self, *args, **kwargs):
          from .frontend.tensorflow2.load import TF2Loader
  
          tf2_loader = TF2Loader(*args, **kwargs)
          return tf2_loader.load()
  
  
  @ConverterRegistry.frontend
  class TorchFrontend:
      name = "pytorch"
  
      def __call__(self, *args, **kwargs):
          from .frontend.torch import load
  
          return load(*args, **kwargs)
  
  
  @ConverterRegistry.backend
  class NNProtoBackend:
      name = "neuralnetwork"
      alias_names = []
  
      def __call__(self, *args, **kwargs):
          from .backend.nn import load
  
          return load(*args, **kwargs)
  
  @ConverterRegistry.backend
  class MILProtoBackend:
      name = "mlprogram"
      alias_names = []
  
      def __call__(self, *args, **kwargs):
          from .backend.mil import load as backend_load
  
          return backend_load(*args, **kwargs)
  
  
  def _reset_conversion_state():
      '''
      Reset any stateful properties/variables that are populated during conversion.
      '''
  
      # Clear the "name_count" dict,
      # which is used to generate unique op names in the mil builder class.
      mb.name_count.clear()
  
      # Clear "k_used_symbols" dict, and the int counter "k_num_internal_syms" that are used to track symbolic names
      global k_used_symbols
      global k_num_internal_syms
      k_used_symbols.clear()
      k_num_internal_syms = 0
  
  @_profile
  def mil_convert(
      model,
      convert_from,
      convert_to,
      compute_units,
      **kwargs
  ):
      """
      Convert model from a specified frontend `convert_from` to a specified
      converter backend `convert_to`.
  
      Parameters
      ----------
      model: TF, PyTorch, or `coremltools.converters.mil.Program`.
          See `coremltools.converters.convert`
  
      convert_from: str
          The value must be one of ['tensorflow', 'tensorflow2',
          'pytorch', 'milinternal'] (aka name of a `ConverterRegistry.frontend`).
  
      compute_units: coremltools.ComputeUnit
          A enum with three possible values:
              - coremltools.ComputeUnit.ALL - use all compute units available, including the
                  neural engine.
              - coremltools.ComputeUnit.CPU_ONLY - limit the model to only use the CPU.
              - coremltools.ComputeUnit.CPU_AND_GPU - use both the CPU and GPU, but not the
                  neural engine.
  
      convert_to: str
         Value must be one of ['neuralnetwork', 'mlprogram', 'milinternal']
         See `coremltools.converters.convert`
  
      Returns
      -------
      model: `coremltools.models.MLModel` or
      `coremltools.converters.mil.Program`
          See `coremltools.converters.convert`
      """
      return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
  
  
  def _mil_convert(
      model,
      convert_from,
      convert_to,
      registry,
      modelClass,
      compute_units,
      **kwargs
  ):
  
      # Map "convert_to" values that correspond to the alias_names, to the actual supported registries
      if convert_to in registry.backend_alias_names:
          msg = "Please use '{}' instead of '{}' with the 'convert_to' argument. The latter will be removed in the future."
          _warnings.warn(msg.format(registry.backend_alias_names[convert_to], convert_to))
          convert_to = registry.backend_alias_names[convert_to]
  
      if convert_to == 'mlprogram':
          # mil_convert_to_proto places weight files inside the weights_dir
          weights_dir = _tempfile.mkdtemp()
          kwargs['weights_dir'] = weights_dir
  
          # To make sure everyone can read and write to this directory (on par with os.mkdir())
          _os.chmod(weights_dir, _stat.S_IRWXU | _stat.S_IRWXG | _stat.S_IRWXO)
  
      proto, mil_program = mil_convert_to_proto(
                              model,
                              convert_from,
                              convert_to,
                              registry,
                              **kwargs
                           )
  
      _reset_conversion_state()
  
      if convert_to == 'milinternal':
          return mil_program # mil program
  
      elif convert_to == 'mlprogram':
          package_path = _create_mlpackage(proto, weights_dir, kwargs.get("package_dir"))
          return modelClass(package_path,
                            is_temp_package=not kwargs.get('package_dir'),
                            mil_program=mil_program,
                            skip_model_load=kwargs.get('skip_model_load', False),
                            compute_units=compute_units)
  
      return modelClass(proto,
                        mil_program=mil_program,
                        skip_model_load=kwargs.get('skip_model_load', False),
                        compute_units=compute_units)
  
  def mil_convert_to_proto(
      model,
      convert_from,
      convert_to,
      converter_registry,
      **kwargs
  ):
      """
      Convert model to proto object.
  
      Parameters
      ----------
      model: See `mil_convert`
  
      convert_from: See `mil_convert`
  
      convert_to: See `mil_convert`
  
      converter_registry: `ConverterRegistry`
        Available frontend and backend converters
  
      Returns
      -------
      model: `coremltools.models.MLModel` or
      `coremltools.converters.mil.Program`
          See `coremltools.converters.convert`
      """
      frontend_converter_type = converter_registry.frontends.get(convert_from.lower())
      if not frontend_converter_type:
          msg = 'Frontend converter "{}" not implemented, must be one of: {}'
          raise NotImplementedError(
              msg.format(convert_from, list(converter_registry.frontends.keys()))
          )
  
      kwargs.setdefault("convert_to", convert_to)
      frontend_converter = frontend_converter_type()
  
      prog = frontend_converter(model, **kwargs)
  
      if convert_to.lower() != "neuralnetwork":
          passes = kwargs.get("transforms", list())
      else:
          # Post training Quantization Passes not available on NN proto backend.
          passes = [p for p in kwargs.get("transforms", list()) if not isinstance(p, AbstractQuantizationPass)]
  
      apply_common_pass_pipeline(prog, passes)
  
      if convert_to == 'milinternal':
          return None, prog # Returns (None, coremltools.converters.mil.Program)
  
      backend_converter_type = converter_registry.backends.get(convert_to.lower())
      if not backend_converter_type:
          msg = 'Backend converter "{}" not implemented, must be one of: {}'
          raise NotImplementedError(
              msg.format(convert_to, list(converter_registry.backends.keys()))
          )
      backend_converter = backend_converter_type()
      out = backend_converter(prog, **kwargs)
  
      return out, prog # Returns (Model_pb2.Model, coremltools.converters.mil.Program)
  