# 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
"""Utilities for handling the benchmark suite.

See https://iree.dev/developers/performance/benchmark-suites/ for how to build
the benchmark suite.
"""

import re

import dataclasses
from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple
from common import benchmark_definition
from common.benchmark_definition import IREE_DRIVERS_INFOS, DriverInfo
from e2e_test_artifacts import iree_artifacts
from e2e_test_framework.definitions import common_definitions, iree_definitions
from e2e_test_framework import serialization


@dataclass
class BenchmarkCase:
    """Represents a benchmark case.

    model_name: the source model, e.g., 'MobileSSD'.
    model_tags: the source model tags, e.g., ['f32'].
    bench_mode: the benchmark mode, e.g., '1-thread,big-core'.
    target_arch: the target CPU/GPU architature.
    driver_info: the IREE driver configuration.
    benchmark_tool_name: the benchmark tool, e.g., 'iree-benchmark-module'.
    run_config: the run config from e2e test framework.
    module_dir: path/URL of the module directory.
    input_uri: URI to find the input npy.
    expected_output_uri: URI to find the expected output npy.
    """

    model_name: str
    model_tags: Sequence[str]
    bench_mode: Sequence[str]
    target_arch: common_definitions.DeviceArchitecture
    driver_info: DriverInfo
    benchmark_tool_name: str
    run_config: iree_definitions.E2EModelRunConfig
    module_dir: benchmark_definition.ResourceLocation
    input_uri: Optional[str] = None
    expected_output_uri: Optional[str] = None
    verify_params: List[str] = dataclasses.field(default_factory=list)
    external_param_urls: List[str] = dataclasses.field(default_factory=list)


# A map from execution config to driver info. This is temporary during migration
# before we can drop the DriverInfo.
EXECUTION_CONFIG_TO_DRIVER_INFO_KEY_MAP: Dict[
    Tuple[iree_definitions.RuntimeDriver, iree_definitions.RuntimeLoader], str
] = {
    (
        iree_definitions.RuntimeDriver.LOCAL_TASK,
        iree_definitions.RuntimeLoader.EMBEDDED_ELF,
    ): "iree-llvm-cpu",
    (
        iree_definitions.RuntimeDriver.LOCAL_SYNC,
        iree_definitions.RuntimeLoader.EMBEDDED_ELF,
    ): "iree-llvm-cpu-sync",
    (
        iree_definitions.RuntimeDriver.LOCAL_TASK,
        iree_definitions.RuntimeLoader.VMVX_MODULE,
    ): "iree-vmvx",
    (
        iree_definitions.RuntimeDriver.LOCAL_SYNC,
        iree_definitions.RuntimeLoader.VMVX_MODULE,
    ): "iree-vmvx-sync",
    (
        iree_definitions.RuntimeDriver.VULKAN,
        iree_definitions.RuntimeLoader.NONE,
    ): "iree-vulkan",
    (
        iree_definitions.RuntimeDriver.CUDA,
        iree_definitions.RuntimeLoader.NONE,
    ): "iree-cuda",
}


class BenchmarkSuite(object):
    """Represents the benchmarks in benchmark suite directory."""

    def __init__(self, benchmark_cases: Sequence[BenchmarkCase]):
        """Construct a benchmark suite.

        Args:
          benchmark_cases: list of benchmark cases.
        """
        self.benchmark_cases = list(benchmark_cases)

    def filter_benchmarks(
        self,
        available_drivers: Optional[Sequence[str]] = None,
        available_loaders: Optional[Sequence[str]] = None,
        target_architectures: Optional[
            Sequence[common_definitions.DeviceArchitecture]
        ] = None,
        driver_filter: Optional[str] = None,
        mode_filter: Optional[str] = None,
        model_name_filter: Optional[str] = None,
    ) -> Sequence[BenchmarkCase]:
        """Filters benchmarks.
        Args:
          available_drivers: list of drivers supported by the tools. None means to
            match any driver.
          available_loaders: list of executable loaders supported by the tools.
            None means to match any loader.
          target_architectures: list of target architectures to be included. None
            means no filter.
          driver_filter: driver filter regex.
          mode_filter: benchmark mode regex.
          model_name_filter: model name regex.
        Returns:
          A list of matched benchmark cases.
        """

        chosen_cases = []
        for benchmark_case in self.benchmark_cases:
            driver_info = benchmark_case.driver_info

            driver_name = driver_info.driver_name
            matched_available_driver = (
                available_drivers is None or driver_name in available_drivers
            )
            matched_driver_filter = (
                driver_filter is None
                or re.match(driver_filter, driver_name) is not None
            )
            matched_driver = matched_available_driver and matched_driver_filter

            matched_loader = (
                not driver_info.loader_name
                or available_loaders is None
                or (driver_info.loader_name in available_loaders)
            )

            if target_architectures is None:
                matched_arch = True
            else:
                matched_arch = benchmark_case.target_arch in target_architectures

            bench_mode = ",".join(benchmark_case.bench_mode)
            matched_mode = (
                mode_filter is None or re.match(mode_filter, bench_mode) is not None
            )

            model_name_with_tags = benchmark_case.model_name
            if len(benchmark_case.model_tags) > 0:
                model_name_with_tags += f"-{','.join(benchmark_case.model_tags)}"
            matched_model_name = (
                model_name_filter is None
                or re.match(model_name_filter, model_name_with_tags) is not None
            )

            if (
                matched_driver
                and matched_loader
                and matched_arch
                and matched_model_name
                and matched_mode
            ):
                chosen_cases.append(benchmark_case)

        return chosen_cases

    @staticmethod
    def load_from_run_configs(
        run_configs: Sequence[iree_definitions.E2EModelRunConfig],
        root_benchmark_dir: benchmark_definition.ResourceLocation,
    ):
        """Loads the benchmarks from the run configs.

        Args:
          run_configs: list of benchmark run configs.
          root_benchmark_dir: path/URL of the root benchmark directory.
        Returns:
          A benchmark suite.
        """

        benchmark_cases = []
        for run_config in run_configs:
            module_gen_config = run_config.module_generation_config
            module_exec_config = run_config.module_execution_config
            target_device_spec = run_config.target_device_spec

            driver_info_key = EXECUTION_CONFIG_TO_DRIVER_INFO_KEY_MAP.get(
                (module_exec_config.driver, module_exec_config.loader)
            )
            if driver_info_key is None:
                raise ValueError(
                    f"Can't map execution config to driver info: {module_exec_config}."
                )
            driver_info = IREE_DRIVERS_INFOS[driver_info_key]

            target_arch = target_device_spec.architecture
            model = module_gen_config.imported_model.model

            module_rel_dir = iree_artifacts.get_module_dir_path(module_gen_config)
            module_dir = root_benchmark_dir / module_rel_dir

            benchmark_case = BenchmarkCase(
                model_name=model.name,
                model_tags=model.tags,
                bench_mode=module_exec_config.tags,
                target_arch=target_arch,
                driver_info=driver_info,
                benchmark_tool_name=run_config.tool.value,
                module_dir=module_dir,
                input_uri=model.input_url,
                expected_output_uri=model.expected_output_url,
                verify_params=model.verify_params,
                external_param_urls=model.external_param_urls,
                run_config=run_config,
            )
            benchmark_cases.append(benchmark_case)

        return BenchmarkSuite(benchmark_cases=benchmark_cases)


def get_run_configs_by_target_and_shard(
    benchmark_groups: Dict, target_device_name: str, shard_index: Optional[int] = None
):
    """Returns a flat list of run_configs from `benchmark_groups`, filtered by the given `target_device_name`.
    If a `shard_index` is given, only the run configs for the given shard are returned, otherwise all the run configs are returned.
    """
    benchmark_group = benchmark_groups.get(target_device_name)
    if benchmark_group is None:
        raise ValueError(
            "Target device '{}' not found in the benchmark config.".format(
                target_device_name
            )
        )

    if shard_index is None:
        # In case no shard index was given we will run ALL benchmarks from ALL shards
        packed_run_configs = [
            shard["run_configs"] for shard in benchmark_group["shards"]
        ]
    else:
        # Otherwise we will only run the benchmarks from the given shard
        benchmark_shard = next(
            (
                shard
                for shard in benchmark_group["shards"]
                if shard["index"] == shard_index
            ),
            None,
        )
        if benchmark_shard is None:
            raise ValueError(
                "Given shard (index={}) not found in the benchmark config group. Available indexes: [{}].".format(
                    shard_index,
                    ", ".join(
                        str(shard["index"]) for shard in benchmark_group["shards"]
                    ),
                )
            )
        packed_run_configs = [benchmark_shard["run_configs"]]

    # When no `shard_index` is given we might have more than one shard to process.
    # We do this by deserializing the `run_config` field from each shard separately
    # and then merge the unpacked flat lists of `E2EModelRunConfig`.
    return [
        run_config
        for packed_run_config in packed_run_configs
        for run_config in serialization.unpack_and_deserialize(
            data=packed_run_config,
            root_type=List[iree_definitions.E2EModelRunConfig],
        )
    ]