# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Utilities to run benchmarks."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numbers
import os
import re
import sys
import time
import types

import six

from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.core.util import test_log_pb2
from tensorflow.python.client import timeline
from tensorflow.python.framework import ops
from tensorflow.python.platform import app
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export


# When a subclass of the Benchmark class is created, it is added to
# the registry automatically
GLOBAL_BENCHMARK_REGISTRY = set()

# Environment variable that determines whether benchmarks are written.
# See also tensorflow/core/util/reporter.h TestReporter::kTestReporterEnv.
TEST_REPORTER_TEST_ENV = "TEST_REPORT_FILE_PREFIX"

# Environment variable that lets the TensorFlow runtime allocate a new
# threadpool for each benchmark.
OVERRIDE_GLOBAL_THREADPOOL = "TF_OVERRIDE_GLOBAL_THREADPOOL"


def _rename_function(f, arg_num, name):
  """Rename the given function's name appears in the stack trace."""
  func_code = six.get_function_code(f)
  if six.PY2:
    new_code = types.CodeType(arg_num, func_code.co_nlocals,
                              func_code.co_stacksize, func_code.co_flags,
                              func_code.co_code, func_code.co_consts,
                              func_code.co_names, func_code.co_varnames,
                              func_code.co_filename, name,
                              func_code.co_firstlineno, func_code.co_lnotab,
                              func_code.co_freevars, func_code.co_cellvars)
  else:
    new_code = types.CodeType(arg_num, 0, func_code.co_nlocals,
                              func_code.co_stacksize, func_code.co_flags,
                              func_code.co_code, func_code.co_consts,
                              func_code.co_names, func_code.co_varnames,
                              func_code.co_filename, name,
                              func_code.co_firstlineno, func_code.co_lnotab,
                              func_code.co_freevars, func_code.co_cellvars)

  return types.FunctionType(new_code, f.__globals__, name, f.__defaults__,
                            f.__closure__)


def _global_report_benchmark(
    name, iters=None, cpu_time=None, wall_time=None,
    throughput=None, extras=None, metrics=None):
  """Method for recording a benchmark directly.

  Args:
    name: The BenchmarkEntry name.
    iters: (optional) How many iterations were run
    cpu_time: (optional) Total cpu time in seconds
    wall_time: (optional) Total wall time in seconds
    throughput: (optional) Throughput (in MB/s)
    extras: (optional) Dict mapping string keys to additional benchmark info.
    metrics: (optional) A list of dict representing metrics generated by the
      benchmark. Each dict should contain keys 'name' and'value'. A dict
      can optionally contain keys 'min_value' and 'max_value'.

  Raises:
    TypeError: if extras is not a dict.
    IOError: if the benchmark output file already exists.
  """
  logging.info("Benchmark [%s] iters: %d, wall_time: %g, cpu_time: %g,"
               "throughput: %g, extras: %s, metrics: %s", name,
               iters if iters is not None else -1,
               wall_time if wall_time is not None else -1,
               cpu_time if cpu_time is not None else -1,
               throughput if throughput is not None else -1,
               str(extras) if extras else "None",
               str(metrics) if metrics else "None")

  entries = test_log_pb2.BenchmarkEntries()
  entry = entries.entry.add()
  entry.name = name
  if iters is not None:
    entry.iters = iters
  if cpu_time is not None:
    entry.cpu_time = cpu_time
  if wall_time is not None:
    entry.wall_time = wall_time
  if throughput is not None:
    entry.throughput = throughput
  if extras is not None:
    if not isinstance(extras, dict):
      raise TypeError("extras must be a dict")
    for (k, v) in extras.items():
      if isinstance(v, numbers.Number):
        entry.extras[k].double_value = v
      else:
        entry.extras[k].string_value = str(v)
  if metrics is not None:
    if not isinstance(metrics, list):
      raise TypeError("metrics must be a list")
    for metric in metrics:
      if "name" not in metric:
        raise TypeError("metric must has a 'name' field")
      if "value" not in metric:
        raise TypeError("metric must has a 'value' field")

      metric_entry = entry.metrics.add()
      metric_entry.name = metric["name"]
      metric_entry.value = metric["value"]
      if "min_value" in metric:
        metric_entry.min_value.value = metric["min_value"]
      if "max_value" in metric:
        metric_entry.max_value.value = metric["max_value"]

  test_env = os.environ.get(TEST_REPORTER_TEST_ENV, None)
  if test_env is None:
    # Reporting was not requested, just print the proto
    print(str(entries))
    return

  serialized_entry = entries.SerializeToString()

  mangled_name = name.replace("/", "__")
  output_path = "%s%s" % (test_env, mangled_name)
  if gfile.Exists(output_path):
    raise IOError("File already exists: %s" % output_path)
  with gfile.GFile(output_path, "wb") as out:
    out.write(serialized_entry)


class _BenchmarkRegistrar(type):
  """The Benchmark class registrar.  Used by abstract Benchmark class."""

  def __new__(mcs, clsname, base, attrs):
    newclass = type.__new__(mcs, clsname, base, attrs)
    if not newclass.is_abstract():
      GLOBAL_BENCHMARK_REGISTRY.add(newclass)
    return newclass


class ParameterizedBenchmark(_BenchmarkRegistrar):
  """Metaclass to generate parameterized benchmarks."""

  def __new__(mcs, clsname, base, attrs):
    param_config_list = attrs["_benchmark_parameters"]