# Copyright 2019 Google LLC
#
# 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
#
#     https://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.
"""Predict extractor for TFLite models."""

import collections
import copy
from typing import Dict, Sequence, Union

from absl import logging
import apache_beam as beam
import numpy as np
import tensorflow as tf
from tensorflow_model_analysis import constants
from tensorflow_model_analysis import types
from tensorflow_model_analysis.extractors import extractor
from tensorflow_model_analysis.proto import config_pb2
from tensorflow_model_analysis.utils import model_util
from tensorflow_model_analysis.utils import util

_TFLITE_PREDICT_EXTRACTOR_STAGE_NAME = 'ExtractTFLitePredictions'


# TODO(b/149981535) Determine if we should merge with RunInference.
@beam.typehints.with_input_types(types.Extracts)
@beam.typehints.with_output_types(types.Extracts)
class _TFLitePredictionDoFn(model_util.BatchReducibleBatchedDoFnWithModels):
  """A DoFn that loads tflite models and predicts."""

  def __init__(self, eval_config: config_pb2.EvalConfig,
               eval_shared_models: Dict[str, types.EvalSharedModel]) -> None:
    super().__init__({k: v.model_loader for k, v in eval_shared_models.items()})
    self._eval_config = eval_config

  def setup(self):
    super().setup()
    self._interpreters = {}

    major, minor, _ = tf.version.VERSION.split('.')
    for model_name, model_contents in self._loaded_models.items():
      # TODO(b/207600661): drop BUILTIN_WITHOUT_DEFAULT_DELEGATES once the issue
      # is fixed.
      if int(major) > 2 or (int(major) == 2 and int(minor) >= 5):
        self._interpreters[model_name] = tf.lite.Interpreter(
            model_content=model_contents.contents,
            experimental_op_resolver_type=tf.lite.experimental.OpResolverType
            .BUILTIN_WITHOUT_DEFAULT_DELEGATES)
      else:
        self._interpreters[model_name] = tf.lite.Interpreter(
            model_content=model_contents.contents)

  def _get_input_name_from_input_detail(self, input_detail):