# Copyright 2020 Google Research. 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.
# ==============================================================================
r"""Inference related utilities."""
import copy
import os
import time
from typing import Text, Dict, Any, Optional
from absl import logging
import numpy as np
import tensorflow as tf

from .. import dataloader
from .. import hparams_config
from .. import utils
from . import efficientdet_keras
from . import label_util
from . import util_keras
from ..visualize import vis_utils


def visualize_image(image,
                    boxes,
                    classes,
                    scores,
                    label_map=None,
                    min_score_thresh=0.01,
                    max_boxes_to_draw=1000,
                    line_thickness=2,
                    **kwargs):
  """Visualizes a given image.

  Args:
    image: a image with shape [H, W, C].
    boxes: a box prediction with shape [N, 4] ordered [ymin, xmin, ymax, xmax].
    classes: a class prediction with shape [N].
    scores: A list of float value with shape [N].
    label_map: a dictionary from class id to name.
    min_score_thresh: minimal score for showing. If claass probability is below
      this threshold, then the object will not show up.
    max_boxes_to_draw: maximum bounding box to draw.
    line_thickness: how thick is the bounding box line.
    **kwargs: extra parameters.

  Returns:
    output_image: an output image with annotated boxes and classes.
  """
  label_map = label_util.get_label_map(label_map or 'coco')
  category_index = {k: {'id': k, 'name': label_map[k]} for k in label_map}
  img = np.array(image)
  vis_utils.visualize_boxes_and_labels_on_image_array(
      img,
      boxes,
      classes,
      scores,
      category_index,
      min_score_thresh=min_score_thresh,
      max_boxes_to_draw=max_boxes_to_draw,
      line_thickness=line_thickness,
      **kwargs)
  return img


class ExportNetwork(tf.Module):

  def __init__(self, model):
    super().__init__()
    self.model = model

  @tf.function
  def __call__(self, imgs):
    return tf.nest.flatten(self.model(imgs, training=False))


class ExportModel(tf.Module):
  """Model to be exported as SavedModel/TFLite format."""

  def __init__(self, model, pre_mode='infer', post_mode='global'):
    super().__init__()
    self.model = model
    self.pre_mode = pre_mode
    self.post_mode = post_mode

  @tf.function
  def __call__(self, imgs):
    return self.model(
        imgs, training=False, pre_mode=self.pre_mode, post_mode=self.post_mode)


class ServingDriver:
  """A driver for serving single or batch images.

  This driver supports serving with image files or arrays, with configurable
  batch size.

  Example 1. Serving streaming image contents:

    driver = inference.ServingDriver(
      'efficientdet-d0', '/tmp/efficientdet-d0', batch_size=1)
    driver.build()
    for m in image_iterator():
      predictions = driver.serve_files([m])
      boxes, scores, classes, _ = tf.nest.map_structure(np.array, predictions)
      driver.visualize(m, boxes[0], scores[0], classes[0])
      # m is the new image with annotated boxes.

  Example 2. Serving batch image contents:

    imgs = []
    for f in ['/tmp/1.jpg', '/tmp/2.jpg']:
      imgs.append(np.array(Image.open(f)))

    driver = inference.ServingDriver(
      'efficientdet-d0', '/tmp/efficientdet-d0', batch_size=len(imgs))
    driver.build()
    predictions = driver.serve(imgs)
    boxes, scores, classes, _ = tf.nest.map_structure(np.array, predictions)
    for i in range(len(imgs)):
      driver.visualize(imgs[i], boxes[i], scores[i], classes[i])

  Example 3: another way is to use SavedModel:

    # step1: export a model.
    driver = inference.ServingDriver('efficientdet-d0', '/tmp/efficientdet-d0')
    driver.build()
    driver.export('/tmp/saved_model_path')

    # step2: Serve a model.
    driver.load(self.saved_model_dir)
    raw_images = []
    for f in tf.io.gfile.glob('/tmp/images/*.jpg'):
      raw_images.append(np.array(PIL.Image.open(f)))
    detections = driver.serve(raw_images)
    boxes, scores, classes, _ = tf.nest.map_structure(np.array, detections)
    for i in range(len(imgs)):
      driver.visualize(imgs[i], boxes[i], scores[i], classes[i])
  """

  def __init__(self,
               model_name: Text,
               ckpt_path: Optional[Text] = None,
               batch_size: int = 1,
               only_network: bool = False,
               model_params: Optional[Dict[Text, Any]] = None,
               debug: bool = False):
    """Initialize the inference driver.

    Args:
      model_name: target model name, such as efficientdet-d0.
      ckpt_path: checkpoint path, such as /tmp/efficientdet-d0/.
      batch_size: batch size for inference.
      only_network: only use the network without pre/post processing.
      model_params: model parameters for overriding the config.
      debug: bool, if true, run in debug mode.
    """
    super().__init__()
    self.model_name = model_name
    self.ckpt_path = ckpt_path
    self.batch_size = batch_size
    self.only_network = only_network
    self.debug = debug

    self.params = hparams_config.get_detection_config(model_name).as_dict()

    if model_params:
      self.params.update(model_params)
    self.params.update(dict(is_training_bn=False))
    self.label_map = self.params.get('label_map', None)

    self._model = None

    mixed_precision = self.params.get('mixed_precision', None)
    precision = utils.get_precision(
        self.params.get('strategy', None), mixed_precision)
    policy = tf.keras.mixed_precision.Policy(precision)
    tf.keras.mixed_precision.set_global_policy(policy)

  @property
  def model(self):
    if not self._model:
      self.build()
    return self._model

  @model.setter
  def model(self, model):
    self._model = model

  def build(self, params_override=None):
    """Build model and restore checkpoints."""
    params = copy.deepcopy(self.params)
    if params_override:
      params.update(params_override)
    config = hparams_config.get_efficientdet_config(self.model_name)
    config.override(params)
    if self.only_network:
      self.model = efficientdet_keras.EfficientDetNet(config=config)
    else:
      self.model = efficientdet_keras.EfficientDetModel(config=config)
    image_size = utils.parse_image_size(params['image_size'])
    self.model.build((self.batch_size, *image_size, 3))
    util_keras.restore_ckpt(self.model, self.ckpt_path,
                            self.params['moving_average_decay'],
                            skip_mismatch=False)
    if self.debug:
      tf.config.run_functions_eagerly(self.debug)

  def visualize(self, image, boxes, classes, scores, **kwargs):
    """Visualize prediction on image."""
    return visualize_image(image, boxes, classes.astype(int), scores,
                           self.label_map, **kwargs)

  def benchmark(self, image_arrays, bm_runs=10, trace_filename=None):
    """Benchmark inference latency/throughput.

    Args:
      image_arrays: a list of images in numpy array format.
      bm_runs: Number of benchmark runs.
      trace_filename: If None, specify the filename for saving trace.
    """
    _, spec = self._get_model_and_spec()

    @tf.function(input_signature=[spec])
    def test_func(image_arrays):
      return self.model(image_arrays)  # pylint: disable=not-callable

    for _ in range(3):  # warmup 3 runs.
      test_func(image_arrays)

    start = time.perf_counter()
    for _ in range(bm_runs):
      test_func(image_arrays)
    end = time.perf_counter()
    inference_time = (end - start) / bm_runs

    print('Per batch inference time: ', inference_time)
    print('FPS: ', self.batch_size / inference_time)

    if trace_filename:
      options = tf.profiler.experimental.ProfilerOptions()
      tf.profiler.experimental.start(trace_filename, options)
      test_func(image_arrays)
      tf.profiler.experimental.stop()

  def serve(self, image_arrays):
    """Serve a list of image arrays.

    Args:
      image_arrays: A list of image content with each image has shape [height,
        width, 3] and uint8 type.

    Returns:
      A list of detections.
    """
    if isinstance(self.model, tf.lite.Interpreter):
      input_details = self.model.get_input_details()
      output_details = self.model.get_output_details()
      self.model.set_tensor(input_details[0]['index'], np.array(image_arrays))
      self.model.invoke()
      return [self.model.get_tensor(x['index']) for x in output_details]
    return self.model(image_arrays)  # pylint: disable=not-callable

  def load(self, saved_model_dir_or_frozen_graph: Text):
    """Load the model using saved model or a frozen graph."""
    # Load saved model if it is a folder.
    if tf.saved_model.contains_saved_model(saved_model_dir_or_frozen_graph):
      self.model = tf.saved_model.load(saved_model_dir_or_frozen_graph)
      return

    if saved_model_dir_or_frozen_graph.endswith('.tflite'):
      self.model = tf.lite.Interpreter(saved_model_dir_or_frozen_graph)
      self.model.allocate_tensors()
      return

    # Load a frozen graph.
    def wrap_frozen_graph(graph_def, inputs, outputs):
      # https://www.tensorflow.org/guide/migrate
      imports_graph_def_fn = lambda: tf.import_graph_def(graph_def, name='')
      wrapped_import = tf.compat.v1.wrap_function(imports_graph_def_fn, [])
      import_graph = wrapped_import.graph
      return wrapped_import.prune(
          tf.nest.map_structure(import_graph.as_graph_element, inputs),
          tf.nest.map_structure(import_graph.as_graph_element, outputs))

    graph_def = tf.Graph().as_graph_def()
    with tf.io.gfile.GFile(saved_model_dir_or_frozen_graph, 'rb') as f:
      graph_def.ParseFromString(f.read())

    self.model = wrap_frozen_graph(
        graph_def,
        inputs='images:0',
        outputs=['Identity:0', 'Identity_1:0', 'Identity_2:0', 'Identity_3:0'])

  def freeze(self, func):
    """Freeze the graph."""
    # pylint: disable=g-import-not-at-top,disable=g-direct-tensorflow-import
    from tensorflow.python.framework.convert_to_constants \
      import convert_variables_to_constants_v2_as_graph
    _, graphdef = convert_variables_to_constants_v2_as_graph(func)
    return graphdef

  def _get_model_and_spec(self, tflite=None):
    """Get model instance and export spec."""
    if self.only_network or tflite:
      image_size = utils.parse_image_size(self.params['image_size'])
      spec = tf.TensorSpec(
          shape=[self.batch_size, *image_size, 3],
          dtype=tf.float32,
          name='images')
      if self.only_network:
        export_model = ExportNetwork(self.model)
      else:
        # If export tflite, we should remove preprocessing since TFLite doesn't
        # support dynamic shape.
        logging.info('Export model without preprocessing.')
        # This section is only used for TFLite, so we use the applicable
        # pre_ & post_ modes.
        export_model = ExportModel(
            self.model, pre_mode=None, post_mode='tflite')
      return export_model, spec
    else:
      spec = tf.TensorSpec(
          shape=[self.batch_size, None, None, 3], dtype=tf.uint8, name='images')
      export_model = ExportModel(self.model)
      return export_model, spec

  def export(self,
             output_dir: Optional[Text] = None,
             tensorrt: Optional[Text] = None,
             tflite: Optional[Text] = None,
             file_pattern: Optional[Text] = None,
             num_calibration_steps: int = 2000):
    """Export a saved model, frozen graph, and potential tflite/tensorrt model.

    Args:
      output_dir: the output folder for saved model.
      tensorrt: If not None, must be {'FP32', 'FP16', 'INT8'}.
      tflite: Type for post-training quantization.
      file_pattern: Glob for tfrecords, e.g. coco/val-*.tfrecord.
      num_calibration_steps: Number of post-training quantization calibration
        steps to run.
    """
    export_model, input_spec = self._get_model_and_spec(tflite)
    image_size = utils.parse_image_size(self.params['image_size'])
    if output_dir:
      tf.saved_model.save(
          export_model,
          output_dir,
          signatures=export_model.__call__.get_concrete_function(input_spec))
      logging.info('Model saved at %s', output_dir)

      # also save freeze pb file.
      graphdef = self.freeze(
          export_model.__call__.get_concrete_function(input_spec))
      proto_path = tf.io.write_graph(
          graphdef, output_dir, self.model_name + '_frozen.pb', as_text=False)
      logging.info('Frozen graph saved at %s', proto_path)

    if tflite:
      shape = (self.batch_size, *image_size, 3)
      input_spec = tf.TensorSpec(
          shape=shape, dtype=input_spec.dtype, name=input_spec.name)
      # from_saved_model supports advanced converter features like op fusing.
      converter = tf.lite.TFLiteConverter.from_saved_model(output_dir)
      if tflite == 'FP32':
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        converter.target_spec.supported_types = [tf.float32]
      elif tflite == 'FP16':
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        converter.target_spec.supported_types = [tf.float16]
      elif tflite == 'INT8':
        # Enables MLIR-based post-training quantization.
        converter.experimental_new_quantizer = True
        if file_pattern:
          config = hparams_config.get_efficientdet_config(self.model_name)
          config.override(self.params)
          ds = dataloader.InputReader(
              file_pattern,
              is_training=False,
              max_instances_per_image=config.max_instances_per_image)(
                  config, batch_size=self.batch_size)

          def representative_dataset_gen():
            for image, _ in ds.take(num_calibration_steps):
              yield [image]
        else:  # Used for debugging, can remove later.
          logging.warn('Use real representative dataset instead of fake ones.')
          num_calibration_steps = 10
          def representative_dataset_gen():  # rewrite this for real data.
            for _ in range(num_calibration_steps):
              yield [tf.ones(shape, dtype=input_spec.dtype)]

        converter.representative_dataset = representative_dataset_gen
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        converter.inference_input_type = tf.uint8
        # TFLite's custom NMS op isn't supported by post-training quant,
        # so we add TFLITE_BUILTINS as well.
        supported_ops = [
            tf.lite.OpsSet.TFLITE_BUILTINS_INT8, tf.lite.OpsSet.TFLITE_BUILTINS
        ]
        converter.target_spec.supported_ops = supported_ops

      else:
        raise ValueError(f'Invalid tflite {tflite}: must be FP32, FP16, INT8.')

      tflite_path = os.path.join(output_dir, tflite.lower() + '.tflite')
      tflite_model = converter.convert()
      tf.io.gfile.GFile(tflite_path, 'wb').write(tflite_model)
      logging.info('TFLite is saved at %s', tflite_path)

    if tensorrt:
      trt_path = os.path.join(output_dir, 'tensorrt_' + tensorrt.lower())
      conversion_params = tf.experimental.tensorrt.ConversionParams(
          max_workspace_size_bytes=(2 << 20),
          maximum_cached_engines=1,
          precision_mode=tensorrt.upper())
      converter = tf.experimental.tensorrt.Converter(
          output_dir, conversion_params=conversion_params)
      converter.convert()
      converter.save(trt_path)
      logging.info('TensorRT model is saved at %s', trt_path)

numpy.array