# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import os
from typing import Optional

from common.benchmark_command import *
from common.benchmark_command_factory import BenchmarkCommandFactory

_DEFAULT_NUM_BENCHMARK_RUNS = 50
_DEFAULT_NUM_THREADS = 1


class TfliteWrapper(TFLiteBenchmarkCommand):
  """Specializes the benchmark command to use TFLite."""

  def __init__(self,
               benchmark_binary: str,
               model_name: str,
               model_path: str,
               input_layer: Optional[str] = None,
               input_shape: Optional[str] = None,
               driver: str = "cpu",
               num_threads: int = _DEFAULT_NUM_THREADS,
               num_runs: int = _DEFAULT_NUM_BENCHMARK_RUNS,
               taskset: Optional[str] = None):
    super().__init__(benchmark_binary,
                     model_name,
                     model_path,
                     num_threads,
                     num_runs,
                     taskset=taskset)
    self.driver = driver
    if input_layer and input_shape:
      self.args.append("--input_layer=%s" % input_layer)
      self.args.append("--input_layer_shape=%s" % input_shape)


class IreeWrapper(IreeBenchmarkCommand):
  """Specializes the benchmark command to use IREE."""

  def __init__(self,
               benchmark_binary: str,
               model_name: str,
               model_path: str,
               function_input: str,
               driver: str = "local-task",
               num_threads: int = _DEFAULT_NUM_THREADS,
               num_runs: int = _DEFAULT_NUM_BENCHMARK_RUNS,
               taskset: Optional[str] = None):
    super().__init__(benchmark_binary,
                     model_name,
                     model_path,
                     num_threads,
                     num_runs,
                     taskset=taskset)
    self.driver = driver
    self.args.append("--entry_function=main")
    self.args.append('--function_input="%s"' % function_input)


class SimpleCommandFactory(BenchmarkCommandFactory):
  """
  Generates `BenchmarkCommand` objects specific to running series of simple models.

  A model is considered simple if its inputs can be generically generated based
  on expected signature only without affecting behavior.
  """

  def __init__(self,
               base_dir: str,
               model_name: str,
               function_input: str,
               input_name: Optional[str] = None,
               input_layer: Optional[str] = None):
    self._model_name = model_name
    self._function_input = function_input
    self._input_name = input_name
    self._input_layer = input_layer
    self._base_dir = base_dir
    self._iree_benchmark_binary_path = os.path.join(base_dir,
                                                    "iree-benchmark-module")
    self._tflite_benchmark_binary_path = os.path.join(base_dir,
                                                      "benchmark_model")
    # Required to be set, but no test data used yet.
    self._tflite_test_data_dir = os.path.join(self._base_dir, "test_data")

  def generate_benchmark_commands(self, device: str,
                                  driver: str) -> list[BenchmarkCommand]:
    if device == "desktop" and driver == "cpu":
      return self._generate_cpu(device)
    elif device == "desktop" and driver == "gpu":
      return self._generate_gpu("cuda")
    elif device == "mobile" and driver == "cpu":
      return self._generate_cpu(device)
    elif device == "mobile" and driver == "gpu":
      return self._generate_gpu("vulkan")
    else:
      print("Warning! Not a valid configuration.")
      return []

  def _generate_cpu(self, device: str):
    commands = []
    # Generate TFLite benchmarks.
    tflite_model_path = os.path.join(self._base_dir, "models", "tflite",
                                     self._model_name + ".tflite")
    tflite = TfliteWrapper(self._tflite_benchmark_binary_path,
                           self._model_name,
                           tflite_model_path,
                           self._input_name,
                           driver="cpu")
    commands.append(tflite)

    # Generate IREE benchmarks.
    driver = "local-task"

    iree_model_path = os.path.join(self._base_dir, "models", "iree", driver,
                                   self._model_name + ".vmfb")
    iree = IreeWrapper(self._iree_benchmark_binary_path,
                       self._model_name,
                       iree_model_path,
                       self._function_input,
                       driver=driver)
    commands.append(iree)

    # Test mmt4d only on mobile.
    if device == "mobile":
      model_mmt4d_name = self._model_name + "_mmt4d"
      iree_mmt4d_model_path = os.path.join(self._base_dir, "models", "iree",
                                           driver, model_mmt4d_name + ".vmfb")
      iree_mmt4d = IreeWrapper(self._iree_benchmark_binary_path,
                               model_mmt4d_name,
                               iree_mmt4d_model_path,
                               self._function_input,
                               driver=driver)
      commands.append(iree_mmt4d)

    return commands

  def _generate_gpu(self, driver: str):
    commands = []
    tflite_model_path = os.path.join(self._base_dir, "models", "tflite",
                                     self._model_name + ".tflite")
    tflite = TfliteWrapper(self._tflite_benchmark_binary_path,
                           self._model_name,
                           tflite_model_path,
                           self._input_name,
                           self._input_layer,
                           driver="gpu")
    tflite.args.append("--gpu_precision_loss_allowed=false")
    commands.append(tflite)

    iree_model_path = os.path.join(self._base_dir, "models", "iree", driver,
                                   self._model_name + ".vmfb")
    iree = IreeWrapper(self._iree_benchmark_binary_path,
                       self._model_name,
                       iree_model_path,
                       self._function_input,
                       driver=driver)
    commands.append(iree)
    return commands