# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""TensorFlow Lite tooling helper functionality."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import warnings
import enum
from six import PY3

from google.protobuf import text_format as _text_format
from google.protobuf.message import DecodeError
from tensorflow.core.framework import graph_pb2 as _graph_pb2
from tensorflow.lite.experimental.examples.lstm.rnn import dynamic_rnn  # pylint: disable=unused-import
from tensorflow.lite.experimental.examples.lstm.rnn_cell import TFLiteLSTMCell  # pylint: disable=unused-import
from tensorflow.lite.experimental.examples.lstm.rnn_cell import TfLiteRNNCell  # pylint: disable=unused-import
from tensorflow.lite.experimental.tensorboard.ops_util import get_potentially_supported_ops  # pylint: disable=unused-import
from tensorflow.lite.python import lite_constants as constants
from tensorflow.lite.python.convert import build_toco_convert_protos  # pylint: disable=unused-import
from tensorflow.lite.python.convert import ConverterError  # pylint: disable=unused-import
from tensorflow.lite.python.convert import OpsSet
from tensorflow.lite.python.convert import toco_convert  # pylint: disable=unused-import
from tensorflow.lite.python.convert import toco_convert_graph_def as _toco_convert_graph_def
from tensorflow.lite.python.convert import toco_convert_impl as _toco_convert_impl
from tensorflow.lite.python.convert import toco_convert_protos  # pylint: disable=unused-import
from tensorflow.lite.python.convert_saved_model import freeze_saved_model as _freeze_saved_model
from tensorflow.lite.python.interpreter import Interpreter  # pylint: disable=unused-import
from tensorflow.lite.python.op_hint import convert_op_hints_to_stubs  # pylint: disable=unused-import
from tensorflow.lite.python.op_hint import OpHint  # pylint: disable=unused-import
from tensorflow.lite.python.optimize import calibrator as _calibrator
from tensorflow.lite.python.util import freeze_graph as _freeze_graph
from tensorflow.lite.python.util import get_grappler_config as _get_grappler_config
from tensorflow.lite.python.util import get_tensor_name as _get_tensor_name
from tensorflow.lite.python.util import get_tensors_from_tensor_names as _get_tensors_from_tensor_names
from tensorflow.lite.python.util import is_frozen_graph as _is_frozen_graph
from tensorflow.lite.python.util import run_graph_optimizations as _run_graph_optimizations
from tensorflow.lite.python.util import set_tensor_shapes as _set_tensor_shapes
from tensorflow.python import keras as _keras
from tensorflow.python.client import session as _session
from tensorflow.python.eager import def_function as _def_function
from tensorflow.python.eager import function as _function
from tensorflow.python.framework import convert_to_constants as _convert_to_constants
from tensorflow.python.framework import dtypes as _dtypes
from tensorflow.python.framework import ops as _ops
from tensorflow.python.framework.errors_impl import NotFoundError as _NotFoundError
from tensorflow.python.framework.importer import import_graph_def as _import_graph_def
from tensorflow.python.lib.io import file_io as _file_io
from tensorflow.python.saved_model import signature_constants as _signature_constants
from tensorflow.python.saved_model import tag_constants as _tag_constants
from tensorflow.python.util import deprecation as _deprecation
from tensorflow.python.util.tf_export import tf_export as _tf_export


@_tf_export("lite.Optimize")
class Optimize(enum.Enum):
  """Enum defining the optimizations to apply when generating tflite graphs.

  Some optimizations may come at the cost of accuracy.
  """

  # Optimize for size.
  #
  # Optimizations that reduce the size of the model.
  # The model size will be reduced. Optimizations can include quantizing the
  # weights of the floating point model.
  OPTIMIZE_FOR_SIZE = "OPTIMIZE_FOR_SIZE"

  # Optimize for latency.
  #
  # Optimizations that reduce the latency of the model.
  # The model latency will be reduced. Optimizations can include quantizing the
  # weights of the floating point model.
  OPTIMIZE_FOR_LATENCY = "OPTIMIZE_FOR_LATENCY"

  def __str__(self):
    return self.value


@_tf_export("lite.RepresentativeDataset")
class RepresentativeDataset(object):
  """Representative dataset to evaluate optimizations.

  A representative dataset that can be used to evaluate optimizations by the
  converter. E.g. converter can use these examples to estimate (min, max) ranges
  by calibrating the model on inputs. This can allow converter to quantize a
  converted floating point model.
  """

  def __init__(self, input_gen, output_gen=None):
    """Creates a representative dataset.

    Args:
      input_gen: an input generator that can be used to generate input samples
        for the model. This must be a callable object that returns an object
        that supports the `iter()` protocol (e.g. a generator function). The
        elements generated must have same type and shape as inputs to the model.
      output_gen: (optional) an output generator that can be used to generate
        output samples for the model. This must be a callable object that
        returns an object that supports the `iter()` protocol (e.g. a generator
        function). The elements generated must have same type and shape as
        outputs to the model. (default None)
    """
    self.input_gen = input_gen
    self.output_gen = output_gen


@_tf_export("lite.TargetSpec")
class TargetSpec(object):
  """Specification of target device.

  Details about target device. Converter optimizes the generated model for
  specific device.

  Attributes:
    supported_ops: Experimental flag, subject to change. Set of OpsSet options
      supported by the device. (default set([OpsSet.TFLITE_BUILTINS]))
  """

  def __init__(self, supported_ops=None):
    if supported_ops is None:
      supported_ops = set([OpsSet.TFLITE_BUILTINS])
    self.supported_ops = supported_ops


@_tf_export("lite.TFLiteConverter", v1=[])
class TFLiteConverterV2(object):
  """Converts a TensorFlow model into TensorFlow Lite model.

  Attributes:
    allow_custom_ops: Boolean indicating whether to allow custom operations.
      When false any unknown operation is an error. When true, custom ops are
      created for any op that is unknown. The developer will need to provide
      these to the TensorFlow Lite runtime with a custom resolver. (default
      False)
    target_spec: Experimental flag, subject to change. Specification of target
      device.
    optimizations: Experimental flag, subject to change, A list of optimizations
      to apply when converting the model. The converter applies the
      optimizations by giving priority to the optimizations specified earlier in
      the list. E.g. `[optimize.OPTIMIZE_FOR_SIZE,
      optimize.OPTIMIZE_FOR_LATENCY]` requires the converter to do both size and
      latency optimizations giving priority to size optimizations over latency
      optimizations.
    representative_dataset: A representative dataset that can be used to
      generate input and output samples for the model. The converter can use the
      dataset to evaluate different optimizations.

  Example usage:

    ```python
    # Converting a GraphDef from a ConcreteFunction.
    converter = lite.TFLiteConverter.from_concrete_function(func)
    tflite_model = converter.convert()
    open("converted_model.tflite", "wb").write(tflite_model)
    ```
  """

  def __init__(self, func):
    """Constructor for TFLiteConverter.

    Args:
      func: TensorFlow ConcreteFunction.
    """
    self._func = func
    self.allow_custom_ops = False
    self.target_spec = TargetSpec()
    self.representative_dataset = None
    self.optimizations = []

  @classmethod
  def from_concrete_function(cls, func):
    """Creates a TFLiteConverter class from a ConcreteFunction.

    Args:
      func: TensorFlow ConcreteFunction.

    Returns:
      TFLiteConverter class.
    """
    if not isinstance(func, _function.ConcreteFunction):
      message = "This function takes in a ConcreteFunction."
      if isinstance(func, _def_function.Function):
        message += (" To get the ConcreteFunction from a Function,"
                    " call from_concrete_function.")
      raise ValueError(message)
    return cls(func)

  def convert(self):
    """Converts a TensorFlow GraphDef based on instance variables.

    Returns:
      The converted data in serialized format.

    Raises:
      ValueError:
        Input shape is not specified.
        None value for dimension in input_tensor.
    """
    frozen_func = _convert_to_constants.convert_variables_to_constants_v2(
        self._func)
    input_tensors = [
        tensor for tensor in frozen_func.inputs
        if tensor.dtype != _dtypes.resource
    ]
    output_tensors = frozen_func.outputs

    # Run a Grappler pass.
    is_only_flex_enabled = set(
        [OpsSet.SELECT_TF_OPS]) == self.target_spec.supported_ops
    config = _get_grappler_config(enable_layout_optimizer=is_only_flex_enabled)
    graph_def = _run_graph_optimizations(
        frozen_func.graph.as_graph_def(),
        input_tensors,
        output_tensors,
        config,
        graph=frozen_func.graph)

    # Checks dimensions in input tensor.
    for tensor in input_tensors:
      # Note that shape_list might be empty for scalar shapes.
      shape_list = tensor.shape.as_list()
      if None in shape_list[1:]:
        raise ValueError(
            "None is only supported in the 1st dimension. Tensor '{0}' has "
            "invalid shape '{1}'.".format(_get_tensor_name(tensor), shape_list))
      elif shape_list and shape_list[0] is None:
        # Set the batch size to 1 if undefined.
        shape = tensor.shape.as_list()
        shape[0] = 1
        tensor.set_shape(shape)

    if self.representative_dataset:
      if not isinstance(self.representative_dataset, RepresentativeDataset):
        raise TypeError("`representative_dataset` must be an instance of "
                        "`RepresentativeDataset`")
      if self.representative_dataset.input_gen is None:
        raise ValueError(
            "Provide an input generator for `representative_dataset`")

    # TODO(shashishekhar): For now use optimizations order is ignored.
    # Both size and latency optimizations decide whether to apply post
    # training optimizations.
    post_training_optimize = bool(
        len(
            set(self.optimizations)
            & set([Optimize.OPTIMIZE_FOR_LATENCY, Optimize.OPTIMIZE_FOR_SIZE])))
    # Do weights only quantization if there is no dataset for calibration.
    weights_only_quantize_flag = (
        post_training_optimize and (self.representative_dataset is None))

    converter_kwargs = {
        "input_format": constants.TENSORFLOW_GRAPHDEF,
        "allow_custom_ops": self.allow_custom_ops,
        "post_training_quantize": weights_only_quantize_flag,
        "target_ops": self.target_spec.supported_ops,
    }

    # Converts model.
    result = _toco_convert_impl(
        input_data=graph_def,
        input_tensors=input_tensors,
        output_tensors=output_tensors,
        **converter_kwargs)

    if self.representative_dataset and post_training_optimize:
      calibrate_quantize = _calibrator.Calibrator(result)
      result = calibrate_quantize.calibrate_and_quantize(
          self.representative_dataset.input_gen)

    return result


@_tf_export(v1=["lite.TFLiteConverter"])
class TFLiteConverter(object):
  """Convert a TensorFlow model into `output_format` using TOCO.

  This is used to convert from a TensorFlow GraphDef or SavedModel into either a
  TFLite FlatBuffer or graph visualization.

  Attributes:

    inference_type: Target data type of real-number arrays in the output file.
      Must be `{tf.float32, tf.uint8}`. (default tf.float32)
    inference_input_type: Target data type of real-number input arrays. Allows
      for a different type for input arrays in the case of quantization.
      Must be `{tf.float32, tf.uint8}`. (default `inference_type`)
    output_format: Output file format. Currently must be `{TFLITE,
      GRAPHVIZ_DOT}`. (default TFLITE)
    quantized_input_stats: Dict of strings representing input tensor names
      mapped to tuple of floats representing the mean and standard deviation
      of the training data (e.g., {"foo" : (0., 1.)}). Only need if
      `inference_input_type` is `QUANTIZED_UINT8`.
      real_input_value = (quantized_input_value - mean_value) / std_dev_value.
      (default {})
    default_ranges_stats: Tuple of integers representing (min, max) range values
      for all arrays without a specified range. Intended for experimenting with
      quantization via "dummy quantization". (default None)
    drop_control_dependency: Boolean indicating whether to drop control
      dependencies silently. This is due to TFLite not supporting control
      dependencies. (default True)
    reorder_across_fake_quant: Boolean indicating whether to reorder FakeQuant
      nodes in unexpected locations. Used when the location of the FakeQuant
      nodes is preventing graph transformations necessary to convert the graph.
      Results in a graph that differs from the quantized training graph,
      potentially causing differing arithmetic behavior. (default False)
    change_concat_input_ranges: Boolean to change behavior of min/max ranges for
      inputs and outputs of the concat operator for quantized models. Changes
      the ranges of concat operator overlap when true. (default False)
    allow_custom_ops: Boolean indicating whether to allow custom operations.
      When false any unknown operation is an error. When true, custom ops are
      created for any op that is unknown. The developer will need to provide
      these to the TensorFlow Lite runtime with a custom resolver.
      (default False)
    post_training_quantize: deprecated, please specify
     `[optimize.OPTIMIZE_FOR_SIZE]` for `optimizations` instead. Boolean
     indicating whether to quantize the weights of the converted float model.
     Model size will be reduced and there will be latency improvements
     (at the cost of accuracy). (default False)
    dump_graphviz_dir: Full filepath of folder to dump the graphs at various
      stages of processing GraphViz .dot files. Preferred over
      --output_format=GRAPHVIZ_DOT in order to keep the requirements of the
      output file. (default None)
    dump_graphviz_video: Boolean indicating whether to dump the graph after
      every graph transformation. (default False)
    target_ops: Experimental flag, subject to change. Set of OpsSet
      options indicating which converter to use.
      (default set([OpsSet.TFLITE_BUILTINS]))
    optimizations: Experimental flag, subject to change, A list of
      optimizations to apply when converting the model. The converter applies
      the optimizations by giving priority to the optimizations specified
      earlier in the list. E.g.
      `[optimize.OPTIMIZE_FOR_SIZE, optimize.OPTIMIZE_FOR_LATENCY]` requires
      the converter to do both size and latency optimizations giving priority
      to size optimizations over latency optimizations.
    representative_dataset: A representative dataset that can be used to