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

Benchmark artifacts should be generated by building the `iree-benchmark-suites`
CMake target, which put them in the following directory structure:

<root-build-dir>/benchmark_suites
└── <benchmark-category> (e.g., TFLite)
    ├── <benchmark-suite> (e.g., MobileBertSquad-fp32)
    │   ├── <benchmark-case> (e.g., iree-vulkan__GPU-Mali-Valhall__kernel-execution)
    │   │   ├── compilation_statistics.json
    │   │   ├── tool
    │   │   └── flagfile
    │   ├── ...
    │   │   ├── compilation_statistics.json
    │   │   ├── tool
    │   │   └── flagfile
    │   └── <benchmark_case>
    │   │   ├── compilation_statistics.json
    │       ├── tool
    │       └── flagfile
    └── vmfb
        ├── <compiled-iree-model>-<sha1>.vmfb
        ├── ...
        └── <compiled-iree-model>-<sha1>.vmfb
"""

import re
import os

from typing import List, Optional, Sequence

from .benchmark_definition import AndroidDeviceInfo, BenchmarkInfo

# All benchmarks' relative path against root build directory.
BENCHMARK_SUITE_REL_PATH = "benchmark_suites"


def compose_info_object(device_info: AndroidDeviceInfo,
                        benchmark_category_dir: str,
                        benchmark_case_dir: str) -> BenchmarkInfo:
  """Creates an BenchmarkInfo object to describe the benchmark.
  Args:
    device_info: an AndroidDeviceInfo object.
    benchmark_category_dir: the directory to a specific benchmark category.
    benchmark_case_dir: a directory containing the benchmark case.
  Returns:
    A BenchmarkInfo object.
  """
  # Extract the model name from the directory path. This uses the relative
  # path under the root model directory. If there are multiple segments,
  # additional ones will be placed in parentheses.
  model_name = os.path.relpath(benchmark_case_dir, benchmark_category_dir)
  # Now we have <model-name>/.../<iree-driver>__<target-arch>__<bench_mode>,
  # Remove the last segment.
  model_name = os.path.dirname(model_name)
  main, rest = os.path.split(model_name)
  if main:
    # Tags coming from directory structure.
    model_name = main
    model_tags = [re.sub(r"\W+", "-", rest)]
  else:
    # Tags coming from the name itself.
    model_name, rest = rest.split("-", 1)
    model_tags = rest.split(",")

  # Extract benchmark info from the directory path following convention:
  #   <iree-driver>__<target-architecture>__<benchmark_mode>
  root_immediate_dir = os.path.basename(benchmark_case_dir)
  iree_driver, target_arch, bench_mode = root_immediate_dir.split("__")

  model_source = os.path.basename(benchmark_category_dir)

  return BenchmarkInfo(model_name=model_name,
                       model_tags=model_tags,
                       model_source=model_source,
                       bench_mode=bench_mode.split(","),
                       runner=iree_driver,
                       device_info=device_info)


def filter_benchmarks_for_category(benchmark_category_dir: str,
                                   cpu_target_arch_filter: str,
                                   gpu_target_arch_filter: str,
                                   driver_filter: Optional[str],
                                   verbose: bool = False) -> Sequence[str]:
  """Filters benchmarks in a specific category for the given device.
  Args:
    benchmark_category_dir: the directory to a specific benchmark category.
    cpu_target_arch_filter: CPU target architecture filter regex.
    gpu_target_arch_filter: GPU target architecture filter regex.
    driver_filter: driver filter regex.
    verbose: whether to print additional debug info.
  Returns:
    A list containing all matched benchmark cases' directories.
  """
  matched_benchmarks = []

  # Go over all benchmarks in the model directory to find those matching the
  # current Android device's CPU/GPU architecture.
  for root, dirs, _ in os.walk(benchmark_category_dir):
    # Take the immediate directory name and try to see if it contains compiled
    # models and flagfiles. This relies on the following directory naming
    # convention:
    #   <iree-driver>__<target-architecture>__<benchmark_mode>
    root_immediate_dir = os.path.basename(root)
    segments = root_immediate_dir.split("__")
    if len(segments) != 3 or not segments[0].startswith("iree-"):
      continue

    iree_driver, target_arch, bench_mode = segments
    iree_driver = iree_driver[len("iree-"):].lower()
    target_arch = target_arch.lower()

    # We can choose this benchmark if it matches the driver and CPU/GPU
    # architecture.
    matched_driver = (driver_filter is None or
                      re.match(driver_filter, iree_driver) is not None)
    matched_arch = (re.match(cpu_target_arch_filter, target_arch) is not None or
                    re.match(gpu_target_arch_filter, target_arch) is not None)
    should_choose = matched_driver and matched_arch
    if should_choose:
      matched_benchmarks.append(root)

    if verbose:
      print(f"dir: {root}")
      print(f"  iree_driver: {iree_driver}")
      print(f"  target_arch: {target_arch}")
      print(f"  bench_mode: {bench_mode}")
      print(f"  chosen: {should_choose}")

  return matched_benchmarks