# 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
"""Provides utilities for benchmarking IREE modules.

Provides convenient methods for invoking IREE's benchmarking tooling from
python. This allows easy benchmarking results from within python.
"""

# pylint: disable=protected-access
# pylint: disable=unused-argument
# pylint: disable=g-explicit-length-test

# TODO(#4131) python>=3.7: Use postponed type annotations.

from collections import namedtuple
from typing import Union
from os import PathLike

from . import VmModule
import numpy
import os
import subprocess
from .dtypes import DTYPE_TO_ABI_TYPE

__all__ = [
    "benchmark_exe",
    "benchmark_module",
]

BenchmarkResult = namedtuple(
    "BenchmarkResult", "benchmark_name time cpu_time iterations user_counters"
)


class BenchmarkToolError(Exception):
    """Benchmark exception that preserves the command line and error output."""

    def __init__(self, message):
        self.message = message
        super().__init__(self.message)


class BenchmarkTimeoutError(Exception):
    """Exception raised if the benchmark is cancelled by the user specified timeout."""

    pass


def benchmark_exe():
    return os.path.join(
        os.path.dirname(__file__), "..", "_runtime_libs", "iree-benchmark-module"
    )


def _build_benchmark_args(
    module: Union[VmModule, PathLike],
    entry_function: str | None = None,
    inputs: list[Union[str, numpy.ndarray]] | None = None,
    **kwargs,
) -> tuple[list[str], bytes | None]:
    args = [benchmark_exe()]

    if isinstance(module, VmModule):
        funcs = [a for a in module.function_names if a != "__init"]
        if entry_function is None:
            if len(funcs) > 1:
                raise ValueError(f"No function specified with multiple options {funcs}")
            entry_function = funcs[0]
        if entry_function not in funcs:
            raise ValueError(
                f"Attempted to benchmark unknown function {entry_function} of options {funcs}"
            )

        flatbuffer = module.stashed_flatbuffer_blob
        args.append("--module=-")
    else:
        flatbuffer = None
        args.append(f"--module={module}")
        if entry_function is None:
            raise ValueError(
                "When specifying the module as a filepath the entry function must be specified."
            )

    args.append(f"--function={entry_function}")

    for k in kwargs:
        v = kwargs[k]
        args.append(f"--{k}={v}")

    for inp in inputs or []:
        if isinstance(inp, str):
            args.append(f"--input={inp}")
            continue
        shape = "x".join([str(d) for d in inp.shape])
        abitype = DTYPE_TO_ABI_TYPE[inp.dtype]
        values = inp.flatten()
        if numpy.all(values[0] == values):
            values = str(values[0])
        else:
            values = ",".join([str(v) for v in values])

        args.append(f"--input={shape}x{abitype}={values}")

    return args, flatbuffer


def benchmark_module(
    module: Union[VmModule, PathLike],
    entry_function: str | None = None,
    inputs: list[Union[str, numpy.ndarray]] | None = None,
    timeout: float | None = None,
    **kwargs,
):
    args, flatbuffer = _build_benchmark_args(
        module=module, entry_function=entry_function, inputs=inputs, **kwargs
    )

    try:
        benchmark_process = subprocess.run(
            args=args,
            input=flatbuffer,
            timeout=timeout,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
    except subprocess.TimeoutExpired:
        raise BenchmarkTimeoutError(f"Benchmark timed out after {timeout} seconds")
    out = benchmark_process.stdout.decode()
    err = benchmark_process.stderr.decode()

    if benchmark_process.returncode != 0:
        raise BenchmarkToolError(f"stderr:\n{err}\nstdout:\n{out}")

    if "INVALID_ARGUMENT;" in err:
        raise ValueError("Invalid inputs specified for benchmarking")

    # In the event benchmarking runs but encounters an internal error,
    # return the internal error instead of benchmark results.
    if "INTERNAL; CUDA driver error" in out:
        raise BenchmarkToolError(out)

    # Grab individual results by line (skip header lines)
    bench_lines = out.split("\n")[3:]
    benchmark_results = []
    for line in bench_lines:
        split = line.split()
        if len(split) == 0:
            continue
        elif len(split) <= 5:
            raise ValueError(
                f"Benchmark standard output is not expected. Line: '{line}' should have at least 6 columns.\nDetails: {args=},\n{out=}"
            )

        benchmark_name = split[0]
        time = " ".join(split[1:3])
        cpu_time = " ".join(split[3:5])
        iterations = split[5]
        user_counters = None
        if len(split) > 5:
            user_counters = split[6]
        benchmark_results.append(
            BenchmarkResult(
                benchmark_name=benchmark_name,
                time=time,
                cpu_time=cpu_time,
                iterations=iterations,
                user_counters=user_counters,
            )
        )

    return benchmark_results
