#!/usr/bin/env python3
# Copyright 2021 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
"""Diffs two local benchmark result JSON files.

Example usage:
  python3 diff_local_benchmarks.py --base=/path/to/base_benchmarks.json
                                   --target=/path/to/target_benchmarks.json
"""

import pathlib
import sys

# Add build_tools python dir to the search path.
sys.path.insert(0, str(pathlib.Path(__file__).parent.with_name("python")))

import argparse

from typing import Optional

from common.benchmark_presentation import *


def get_benchmark_result_markdown(
    base_benchmark_file: Optional[pathlib.Path],
    target_benchmark_file: Optional[pathlib.Path],
    base_compile_stats_file: Optional[pathlib.Path],
    target_compile_stats_file: Optional[pathlib.Path],
    verbose: bool = False) -> str:
  """Gets the full markdown summary of all benchmarks in files."""
  base_benchmarks = {}
  target_benchmarks = {}
  base_compilation_metrics = {}
  target_compilation_metrics = {}
  if base_benchmark_file and target_benchmark_file:
    base_benchmarks = aggregate_all_benchmarks([base_benchmark_file])
    target_benchmarks = aggregate_all_benchmarks([target_benchmark_file])
  if base_compile_stats_file and target_compile_stats_file:
    base_compilation_metrics = collect_all_compilation_metrics(
        [base_compile_stats_file])
    target_compilation_metrics = collect_all_compilation_metrics(
        [target_compile_stats_file])

  # Update the target benchmarks with their corresponding base numbers.
  for bench in base_benchmarks:
    if bench in target_benchmarks:
      target_benchmarks[bench].base_mean_time = base_benchmarks[bench].mean_time

  for target_name, base_metrics in base_compilation_metrics.items():
    updated_metrics = base_metrics
    for mapper in COMPILATION_METRICS_TO_TABLE_MAPPERS:
      metric_key = mapper.get_series_name(target_name)
      base_value, _ = mapper.get_current_and_base_value(base_metrics)
      updated_metrics = mapper.update_base_value(updated_metrics, base_value)
    target_compilation_metrics[target_name] = updated_metrics

  # Compose the full benchmark tables.
  full_table = [md.header("Full Benchmark Summary", 2)]
  full_table.append(categorize_benchmarks_into_tables(target_benchmarks))

  # Compose the full compilation metrics tables.
  full_table.append(
      categorize_compilation_metrics_into_tables(target_compilation_metrics))

  return "\n\n".join(full_table)


def parse_arguments():
  """Parses command-line options."""

  def check_file_path(path):
    path = pathlib.Path(path)
    if path.is_file():
      return path
    else:
      raise ValueError(path)

  parser = argparse.ArgumentParser()
  parser.add_argument("--base",
                      type=check_file_path,
                      help="Base benchmark results")
  parser.add_argument("--target",
                      type=check_file_path,
                      help="Target benchmark results")
  parser.add_argument("--base-compile-stats",
                      type=check_file_path,
                      help="Base compilation statistics")
  parser.add_argument("--target-compile-stats",
                      type=check_file_path,
                      help="Target compilation statistics")
  parser.add_argument("--verbose",
                      action="store_true",
                      help="Print internal information during execution")
  args = parser.parse_args()

  return args


if __name__ == "__main__":
  args = parse_arguments()
  if args.base or args.target:
    if not args.base or not args.target:
      raise ValueError("--base and --target must be used together.")
  if args.base_compile_stats or args.target_compile_stats:
    if not args.base_compile_stats or not args.target_compile_stats:
      raise ValueError("--base-compile-stats and --target-compile-stats must "
                       "be used together.")

  print(
      get_benchmark_result_markdown(args.base,
                                    args.target,
                                    args.base_compile_stats,
                                    args.target_compile_stats,
                                    verbose=args.verbose))