# Copyright 2020 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.
"""Model specification for object detection."""

import collections
import os
import tempfile
from typing import Optional, Tuple, Dict

from absl import logging
import tensorflow as tf
from tensorflow_examples.lite.model_maker.core import compat
from tensorflow_examples.lite.model_maker.core.data_util import object_detector_dataloader
from tensorflow_examples.lite.model_maker.core.export_format import QuantizationType
from tensorflow_examples.lite.model_maker.core.task import configs
from tensorflow_examples.lite.model_maker.core.task.model_spec import util

from tensorflow_examples.lite.model_maker.third_party.efficientdet import coco_metric
from tensorflow_examples.lite.model_maker.third_party.efficientdet import hparams_config
from tensorflow_examples.lite.model_maker.third_party.efficientdet import utils
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import efficientdet_keras
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import eval_tflite
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import inference
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import label_util
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import postprocess
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import train
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import train_lib
from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import util_keras


# Number of calibration steps for full integer quantization. 500 steps are
# enough to get a reasonable post-quantization result.
_NUM_CALIBRATION_STEPS = 500


def _get_ordered_label_map(
    label_map: Optional[Dict[int, str]]) -> Optional[Dict[int, str]]:
  """Gets label_map as an OrderedDict instance with ids sorted."""
  if not label_map:
    return label_map
  ordered_label_map = collections.OrderedDict()
  for idx in sorted(label_map.keys()):
    ordered_label_map[idx] = label_map[idx]
  return ordered_label_map


class EfficientDetModelSpec(object):
  """A specification of the EfficientDet model."""

  compat_tf_versions = compat.get_compat_tf_versions(2)

  def __init__(self,
               model_name: str,
               uri: str,
               hparams: str = '',
               model_dir: Optional[str] = None,
               epochs: int = 50,
               batch_size: int = 64,
               steps_per_execution: int = 1,
               moving_average_decay: int = 0,
               var_freeze_expr: str = '(efficientnet|fpn_cells|resample_p6)',
               tflite_max_detections: int = 25,
               strategy: Optional[str] = None,
               tpu: Optional[str] = None,
               gcp_project: Optional[str] = None,
               tpu_zone: Optional[str] = None,
               use_xla: bool = False,
               profile: bool = False,
               debug: bool = False,
               tf_random_seed: int = 111111) -> None:
    """Initialze an instance with model paramaters.

    Args:
      model_name: Model name.
      uri: TF-Hub path/url to EfficientDet module.
      hparams: Hyperparameters used to overwrite default configuration. Can be
        1) Dict, contains parameter names and values; 2) String, Comma separated
        k=v pairs of hyperparameters; 3) String, yaml filename which's a module
        containing attributes to use as hyperparameters.
      model_dir: The location to save the model checkpoint files.
      epochs: Default training epochs.
      batch_size: Training & Evaluation batch size.
      steps_per_execution: Number of steps per training execution.
      moving_average_decay: Float. The decay to use for maintaining moving
        averages of the trained parameters.
      var_freeze_expr: Expression to freeze variables.
      tflite_max_detections: The max number of output detections in the TFLite
        model.
      strategy:  A string specifying which distribution strategy to use.
        Accepted values are 'tpu', 'gpus', None. tpu' means to use TPUStrategy.
        'gpus' mean to use MirroredStrategy for multi-gpus. If None, use TF
        default with OneDeviceStrategy.
      tpu: The Cloud TPU to use for training. This should be either the name
        used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470
          url.
      gcp_project: Project name for the Cloud TPU-enabled project. If not
        specified, we will attempt to automatically detect the GCE project from
        metadata.
      tpu_zone: GCE zone where the Cloud TPU is located in. If not specified, we
        will attempt to automatically detect the GCE project from metadata.
      use_xla: Use XLA even if strategy is not tpu. If strategy is tpu, always
        use XLA, and this flag has no effect.
      profile: Enable profile mode.
      debug: Enable debug mode.
      tf_random_seed: Fixed random seed for deterministic execution across runs
        for debugging.
    """
    self.model_name = model_name
    self.uri = uri
    self.batch_size = batch_size
    config = hparams_config.get_efficientdet_config(model_name)
    config.override(hparams)
    config.image_size = utils.parse_image_size(config.image_size)
    config.var_freeze_expr = var_freeze_expr
    config.moving_average_decay = moving_average_decay
    config.tflite_max_detections = tflite_max_detections
    if epochs:
      config.num_epochs = epochs

    if use_xla and strategy != 'tpu':
      tf.config.optimizer.set_jit(True)
      for gpu in tf.config.list_physical_devices('GPU'):
        tf.config.experimental.set_memory_growth(gpu, True)

    if debug:
      tf.config.experimental_run_functions_eagerly(True)
      tf.debugging.set_log_device_placement(True)
      os.environ['TF_DETERMINISTIC_OPS'] = '1'
      tf.random.set_seed(tf_random_seed)
      logging.set_verbosity(logging.DEBUG)

    if strategy == 'tpu':
      tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
          tpu, zone=tpu_zone, project=gcp_project)
      tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
      tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)
      ds_strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
      logging.info('All devices: %s', tf.config.list_logical_devices('TPU'))
      tf.config.set_soft_device_placement(True)
    elif strategy == 'gpus':
      ds_strategy = tf.distribute.MirroredStrategy()
      logging.info('All devices: %s', tf.config.list_physical_devices('GPU'))
    else:
      if tf.config.list_physical_devices('GPU'):
        ds_strategy = tf.distribute.OneDeviceStrategy('device:GPU:0')
      else:
        ds_strategy = tf.distribute.OneDeviceStrategy('device:CPU:0')

    self.ds_strategy = ds_strategy

    if model_dir is None:
      model_dir = tempfile.mkdtemp()
    params = dict(
        profile=profile,
        model_name=model_name,
        steps_per_execution=steps_per_execution,
        model_dir=model_dir,
        strategy=strategy,
        batch_size=batch_size,
        tf_random_seed=tf_random_seed,
        debug=debug)
    config.override(params, True)
    self.config = config

    # set mixed precision policy by keras api.
    precision = utils.get_precision(config.strategy, config.mixed_precision)
    policy = tf.keras.mixed_precision.experimental.Policy(precision)
    tf.keras.mixed_precision.experimental.set_policy(policy)

  def create_model(self) -> tf.keras.Model:
    """Creates the EfficientDet model."""
    return train_lib.EfficientDetNetTrainHub(
        config=self.config, hub_module_url=self.uri)

  def train(self,
            model: tf.keras.Model,
            train_dataset: tf.data.Dataset,
            steps_per_epoch: int,
            val_dataset: Optional[tf.data.Dataset],
            validation_steps: int,
            epochs: Optional[int] = None,
            batch_size: Optional[int] = None,
            val_json_file: Optional[str] = None) -> tf.keras.Model:
    """Run EfficientDet training."""
    config = self.config
    if not epochs:
      epochs = config.num_epochs

    if not batch_size:
      batch_size = config.batch_size

    config.update(
        dict(
            steps_per_epoch=steps_per_epoch,
            eval_samples=batch_size * validation_steps,
            val_json_file=val_json_file,
            batch_size=batch_size))
    train.setup_model(model, config)
    train.init_experimental(config)
    model.fit(
        train_dataset,
        epochs=epochs,
        steps_per_epoch=steps_per_epoch,
        callbacks=train_lib.get_callbacks(config.as_dict(), val_dataset),
        validation_data=val_dataset,
        validation_steps=validation_steps)
    return model

  def _get_evaluator_and_label_map(
      self, json_file: str
  ) -> Tuple[coco_metric.EvaluationMetric, Optional[Dict[int, str]]]:
    """Gets evaluator and label_map for evaluation."""
    label_map = label_util.get_label_map(self.config.label_map)
    # Sorts label_map.keys since pycocotools.cocoeval uses sorted catIds
    # (category ids) in COCOeval class.
    label_map = _get_ordered_label_map(label_map)

    evaluator = coco_metric.EvaluationMetric(
        filename=json_file, label_map=label_map)

    evaluator.reset_states()
    return evaluator, label_map

  def _get_metric_dict(self, evaluator: coco_metric.EvaluationMetric,
                       label_map: collections.OrderedDict) -> Dict[str, float]:
    """Gets the metric dict for evaluation."""
    metrics = evaluator.result()
    metric_dict = {}
    for i, name in enumerate(evaluator.metric_names):
      metric_dict[name] = metrics[i]

    if label_map:
      for i, cid in enumerate(label_map.keys()):
        name = 'AP_/%s' % label_map[cid]
        metric_dict[name] = metrics[i + len(evaluator.metric_names)]
    return metric_dict

  def evaluate(self,
               model: tf.keras.Model,
               dataset: tf.data.Dataset,
               steps: int,
               json_file: str = None) -> Dict[str, float]:
    """Evaluate the EfficientDet keras model.

    Args:
      model: The keras model to be evaluated.
      dataset: tf.data.Dataset used for evaluation.
      steps: Number of steps to evaluate the model.
      json_file: JSON with COCO data format containing golden bounding boxes.
        Used for validation. If None, use the ground truth from the dataloader.
        Refer to
        https://towardsdatascience.com/coco-data-format-for-object-detection-a4c5eaf518c5
          for the description of COCO data format.

    Returns:
      A dict contains AP metrics.
    """
    evaluator, label_map = self._get_evaluator_and_label_map(json_file)
    dataset = dataset.take(steps)

    @tf.function
    def _get_detections(images, labels):
      cls_outputs, box_outputs = model(images, training=False)
      detections = postprocess.generate_detections(self.config, cls_outputs,
                                                   box_outputs,
                                                   labels['image_scales'],
                                                   labels['source_ids'])
      tf.numpy_function(evaluator.update_state, [
          labels['groundtruth_data'],
          postprocess.transform_detections(detections)
      ], [])

    dataset = self.ds_strategy.experimental_distribute_dataset(dataset)
    progbar = tf.keras.utils.Progbar(steps)
    for i, (images, labels) in enumerate(dataset):
      self.ds_strategy.run(_get_detections, (images, labels))
      progbar.update(i)

    metric_dict = self._get_metric_dict(evaluator, label_map)
    return metric_dict

  def evaluate_tflite(self,
                      tflite_filepath: str,
                      dataset: tf.data.Dataset,
                      steps: int,
                      json_file: str = None) -> Dict[str, float]:
    """Evaluate the EfficientDet TFLite model.

    Args:
      tflite_filepath: File path to the TFLite model.
      dataset: tf.data.Dataset used for evaluation.
      steps: Number of steps to evaluate the model.
      json_file: JSON with COCO data format containing golden bounding boxes.
        Used for validation. If None, use the ground truth from the dataloader.
        Refer to
        https://towardsdatascience.com/coco-data-format-for-object-detection-a4c5eaf518c5
          for the description of COCO data format.

    Returns:
      A dict contains AP metrics.
    """
    # TODO(b/182441458): Use the task library for evaluation instead once it
    # supports python interface.
    evaluator, label_map = self._get_evaluator_and_label_map(json_file)
    dataset = dataset.take(steps)

    lite_runner = eval_tflite.LiteRunner(tflite_filepath, only_network=False)
    progbar = tf.keras.utils.Progbar(steps)
    for i, (images, labels) in enumerate(dataset):
      # Get the output result after post-processing NMS op.
      nms_boxes, nms_classes, nms_scores, _ = lite_runner.run(images)

      # CLASS_OFFSET is used since label_id for `background` is 0 in label_map
      # while it's not actually included the model. We don't need to add the
      # offset in the Android application.
      nms_classes += postprocess.CLASS_OFFSET

      height, width = utils.parse_image_size(self.config.image_size)
      normalize_factor = tf.constant([height, width, height, width],
                                     dtype=tf.float32)
      nms_boxes *= normalize_factor
      if labels['image_scales'] is not None:
        scales = tf.expand_dims(tf.expand_dims(labels['image_scales'], -1), -1)
        nms_boxes = nms_boxes * tf.cast(scales, nms_boxes.dtype)
      detections = postprocess.generate_detections_from_nms_output(
          nms_boxes, nms_classes, nms_scores, labels['source_ids'])

      detections = postprocess.transform_detections(detections)
      evaluator.update_state(labels['groundtruth_data'].numpy(),
                             detections.numpy())
      progbar.update(i)

    metric_dict = self._get_metric_dict(evaluator, label_map)
    return metric_dict

  def export_saved_model(self,
                         saved_model_dir: str,
                         batch_size: Optional[int] = None,
                         pre_mode: Optional[str] = 'infer',
                         post_mode: Optional[str] = 'global') -> None:
    """Saves the model to Tensorflow SavedModel.

    Args:
      saved_model_dir: Folder path for saved model.
      batch_size: Batch size to be saved in saved_model.
      pre_mode: Pre-processing Mode in ExportModel, must be {None, 'infer'}.
      post_mode: Post-processing Mode in ExportModel, must be {None, 'global',
        'per_class'}.
    """
    # Create EfficientDetModel with latest checkpoint.
    config = self.config
    tf.keras.backend.clear_session()
    model = efficientdet_keras.EfficientDetModel(config=config)
    model.build((batch_size, *config.image_size, 3))
    if config.model_dir:
      util_keras.restore_ckpt(
          model,
          config.model_dir,
          config['moving_average_decay'],
          skip_mismatch=False)
    else:
      # EfficientDetModel is random initialized without restoring the
      # checkpoint. This is mainly used in object_detector_test and shouldn't be
      #  used if we want to export trained model.
      tf.compat.v1.logging.warn('Need to restore the checkpoint for '
                                'EfficientDet.')
    # Gets tf.TensorSpec.
    if pre_mode is None:
      # Input is the preprocessed image that's already resized to a certain
      # input shape.
      input_spec = tf.TensorSpec(
          shape=[batch_size, *config.image_size, 3],
          dtype=tf.float32,
          name='images')
    else:
      # Input is that raw image that can be in any input shape,
      input_spec = tf.TensorSpec(
          shape=[batch_size, None, None, 3], dtype=tf.uint8, name='images')

    export_model = inference.ExportModel(
        model, pre_mode=pre_mode, post_mode=post_mode)
    tf.saved_model.save(
        export_model,
        saved_model_dir,
        signatures=export_model.__call__.get_concrete_function(input_spec))

  def export_tflite(
      self,
      tflite_filepath: str,
      quantization_type: Optional[QuantizationType] = None,
      representative_dataset: Optional[
          object_detector_dataloader.DataLoader] = None,
      quantization_config: Optional[configs.QuantizationConfig] = None) -> None:
    """Converts the retrained model to tflite format and saves it.

    The exported TFLite model has the following inputs & outputs:
    One input:
      image: a float32 tensor of shape[1, height, width, 3] containing the
        normalized input image. `self.config.image_size` is [height, width].

    Four Outputs:
      detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box
        locations.
      detection_classes: a float32 tensor of shape [1, num_boxes] with class
        indices.
      detection_scores: a float32 tensor of shape [1, num_boxes] with class
        scores.
      num_boxes: a float32 tensor of size 1 containing the number of detected
        boxes.

    Args:
      tflite_filepath: File path to save tflite model.
      quantization_type: Enum, type of post-training quantization. Accepted
        values are `INT8`, `FP16`, `FP32`, `DYNAMIC`. `FP16` means float16
        quantization with 2x smaller, optimized for GPU. `INT8` means full
        integer quantization with 4x smaller, 3x+ speedup, optimized for Edge
        TPU. 'DYNAMIC' means dynamic range quantization with	4x smaller, 2x-3x
        speedup. `FP32` mean exporting float model without quantization. Please
        refer to
        https://www.tensorflow.org/lite/performance/post_training_quantization
        for more detailed about different techniques for post-training
        quantization.
      representative_dataset: Representative dataset for full integer
        quantization. Used when `quantization_type=INT8`.
      quantization_config: Configuration for post-training quantization.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
      self.export_saved_model(
          temp_dir, batch_size=1, pre_mode=None, post_mode='tflite')
      converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir)
      if quantization_type == QuantizationType.INT8:
        # Enables MLIR-based post-training quantization.
        converter.experimental_new_quantizer = True

        if representative_dataset is None:
          raise ValueError('`representative_data` must be set when '
                           '`quantization_type=QuantizationType.INT8.')

        def representative_dataset_gen():
          for image, _ in representative_dataset.take(_NUM_CALIBRATION_STEPS):
            yield [image]

        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
      elif quantization_type == QuantizationType.FP16:
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        converter.target_spec.supported_types = [tf.float16]
      elif quantization_type == QuantizationType.DYNAMIC:
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
      elif quantization_type != QuantizationType.FP32:
        raise ValueError('Unsupported `quantization_type`: %s' %
                         str(quantization_type))
      elif quantization_config:
        converter = quantization_config.get_converter_with_quantization(
            converter, model_spec=self)

        # TFLITE_BUILTINS is needed for TFLite's custom NMS op for integer only
        # quantization.
        if tf.lite.OpsSet.TFLITE_BUILTINS not in converter.target_spec.supported_ops:
          converter.target_spec.supported_ops += [
              tf.lite.OpsSet.TFLITE_BUILTINS
          ]

      tflite_model = converter.convert()

      with tf.io.gfile.GFile(tflite_filepath, 'wb') as f:
        f.write(tflite_model)


def efficientdet_lite0_spec(**kwargs) -> EfficientDetModelSpec:
  args = util.dict_with_default(
      default_dict=dict(
          model_name='efficientdet-lite0',
          uri='https://tfhub.dev/tensorflow/efficientdet/lite0/feature-vector/1'
      ),
      **kwargs)
  return EfficientDetModelSpec(**args)


def efficientdet_lite1_spec(**kwargs) -> EfficientDetModelSpec:
  args = util.dict_with_default(
      default_dict=dict(
          model_name='efficientdet-lite1',
          uri='https://tfhub.dev/tensorflow/efficientdet/lite1/feature-vector/1'
      ),
      **kwargs)
  return EfficientDetModelSpec(**args)


def efficientdet_lite2_spec(**kwargs) -> EfficientDetModelSpec:
  args = util.dict_with_default(
      default_dict=dict(
          model_name='efficientdet-lite2',
          uri='https://tfhub.dev/tensorflow/efficientdet/lite2/feature-vector/1'
      ),
      **kwargs)
  return EfficientDetModelSpec(**args)


def efficientdet_lite3_spec(**kwargs) -> EfficientDetModelSpec:
  args = util.dict_with_default(
      default_dict=dict(
          model_name='efficientdet-lite3',
          uri='https://tfhub.dev/tensorflow/efficientdet/lite3/feature-vector/1'
      ),
      **kwargs)
  return EfficientDetModelSpec(**args)


def efficientdet_lite4_spec(**kwargs) -> EfficientDetModelSpec:
  args = util.dict_with_default(
      default_dict=dict(
          model_name='efficientdet-lite4',
          uri='https://tfhub.dev/tensorflow/efficientdet/lite4/feature-vector/1'
      ),
      **kwargs)
  return EfficientDetModelSpec(**args)
