# 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.

import math
import unittest

from . import test_utils
from edgetpu.basic import edgetpu_utils
from edgetpu.basic.basic_engine import BasicEngine


class TestBasicEnginePythonAPI(unittest.TestCase):

  def testDebugInfo(self):
    engine = BasicEngine(
        test_utils.TestDataPath('mobilenet_v1_1.0_224_quant.tflite'))
    # Check model's input format.
    input_tensor_shape = engine.get_input_tensor_shape()
    self.assertListEqual([1, 224, 224, 3], input_tensor_shape.tolist())
    self.assertEqual(224 * 224 * 3, engine.required_input_array_size())

    # Check model's output.
    output_tensors_sizes = engine.get_all_output_tensors_sizes()
    self.assertListEqual([1001], output_tensors_sizes.tolist())
    self.assertEqual(1, engine.get_num_of_output_tensors())
    self.assertEqual(1001, engine.get_output_tensor_size(0))
    self.assertEqual(1001, engine.total_output_array_size())

    # Check SSD model.
    ssd_engine = BasicEngine(
        test_utils.TestDataPath(
            'mobilenet_ssd_v1_coco_quant_postprocess.tflite'))
    # Check model's input format.
    input_tensor_shape = ssd_engine.get_input_tensor_shape()
    self.assertListEqual([1, 300, 300, 3], input_tensor_shape.tolist())
    self.assertEqual(300 * 300 * 3, ssd_engine.required_input_array_size())

    # Check model's output.
    output_tensors_sizes = ssd_engine.get_all_output_tensors_sizes()
    self.assertListEqual([40, 10, 10, 1], output_tensors_sizes.tolist())
    self.assertEqual(4, ssd_engine.get_num_of_output_tensors())
    self.assertEqual(40, ssd_engine.get_output_tensor_size(0))
    self.assertEqual(10, ssd_engine.get_output_tensor_size(1))
    self.assertEqual(10, ssd_engine.get_output_tensor_size(2))
    self.assertEqual(1, ssd_engine.get_output_tensor_size(3))
    self.assertEqual(61, ssd_engine.total_output_array_size())

  def testRunInference(self):
    for model in test_utils.GetModelList():
      print('Testing model :', model)
      engine = BasicEngine(test_utils.TestDataPath(model))
      input_data = test_utils.GenerateRandomInput(
          1, engine.required_input_array_size())
      latency, ret = engine.RunInference(input_data)
      self.assertEqual(ret.size, engine.total_output_array_size())
      # Check debugging functions.
      self.assertLess(math.fabs(engine.get_inference_time() - latency), 0.001)
      raw_output = engine.get_raw_output()
      self.assertEqual(ret.size, raw_output.size)
      for i in range(ret.size):
        if math.isnan(ret[i]) and math.isnan(raw_output[i]):
          continue
        self.assertLess(math.fabs(ret[i] - raw_output[i]), 0.001)

  def testDevicePath(self):
    all_edgetpu_paths = edgetpu_utils.ListEdgeTpuPaths(
        edgetpu_utils.EDGE_TPU_STATE_NONE)
    engine = BasicEngine(
        test_utils.TestDataPath('mobilenet_v1_1.0_224_quant.tflite'),
        all_edgetpu_paths[0])
    self.assertEqual(engine.device_path(), all_edgetpu_paths[0])

if __name__ == '__main__':
  unittest.main()