# 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("--function=main")
        self.args.append('--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)

        tflite_noxnn = TfliteWrapper(
            self._tflite_benchmark_binary_path,
            self._model_name + "_noxnn",
            tflite_model_path,
            self._input_name,
            driver="cpu",
        )
        tflite_noxnn.args.append("--use_xnnpack=false")
        commands.append(tflite_noxnn)

        # Generate IREE benchmarks.
        driver = "local-task"
        backend = "llvm-cpu"

        iree_model_path = os.path.join(
            self._base_dir, "models", "iree", backend, 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)

        model_padfuse_name = self._model_name + "_padfuse"
        iree_padfuse_model_path = os.path.join(
            self._base_dir, "models", "iree", backend, model_padfuse_name + ".vmfb"
        )
        iree_padfuse = IreeWrapper(
            self._iree_benchmark_binary_path,
            model_padfuse_name,
            iree_padfuse_model_path,
            self._function_input,
            driver=driver,
        )
        commands.append(iree_padfuse)

        # 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", backend, 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)

            model_im2col_mmt4d_name = self._model_name + "_im2col_mmt4d"
            iree_im2col_mmt4d_model_path = os.path.join(
                self._base_dir,
                "models",
                "iree",
                backend,
                model_im2col_mmt4d_name + ".vmfb",
            )
            iree_im2col_mmt4d = IreeWrapper(
                self._iree_benchmark_binary_path,
                model_im2col_mmt4d_name,
                iree_im2col_mmt4d_model_path,
                self._function_input,
                driver=driver,
            )
            commands.append(iree_im2col_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)

        tflite_noxnn = TfliteWrapper(
            self._tflite_benchmark_binary_path,
            self._model_name + "_noxnn",
            tflite_model_path,
            self._input_name,
            self._input_layer,
            driver="gpu",
        )
        tflite.args.append("--use_xnnpack=false")
        commands.append(tflite_noxnn)

        tflite_fp16 = TfliteWrapper(
            self._tflite_benchmark_binary_path,
            self._model_name + "_fp16",
            tflite_model_path,
            self._input_name,
            self._input_layer,
            driver="gpu",
        )
        tflite.args.append("--gpu_precision_loss_allowed=true")
        commands.append(tflite_fp16)

        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)

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

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