#!/usr/bin/env python3

# Copyright 2021 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import logging
import os
import pandas as pd
import numpy

from scipy import stats as scipy_stats


def get_diamond_string(diamond_count: int):
  if diamond_count == 0:
    return "~"
  elif diamond_count == 1:
    return "◆"
  elif diamond_count == 2:
    return "◆◆"
  elif diamond_count == 3:
    return "◆◆◆"
  elif diamond_count == 4:
    return "◆◆◆◆"


def get_diamonds_count(significance: pd.DataFrame):
  """
  This function emulates the "diamond" significance representation
  that is familiar to UMA users.
  """

  assert (not (significance > 1).any().any())
  assert (not (significance < 0).any().any())

  # Avoid log10(0) which is undefined.
  significance = numpy.clip(significance, 0, 0.999999)

  # scipy_stats.norm.cdf(1.96) = 0.975 and we're interested in 2 tail
  # test. 1.96 gives a 0.05 p-value. Multiply by 2 here to correct.
  p_value = (1 - significance) * 2

  # floor() to avoid exaggerating results and to round.
  # absolute() to make the result positive.
  log_p_value = numpy.floor(numpy.absolute(numpy.log10(p_value)))

  # Clip because 4 diamond is the max no matter the p-value.
  return numpy.clip(log_p_value, 0, 4)


def compute_mean_and_stderr(summary_path: str):
  df = pd.read_csv(summary_path)

  # skipna because no line has as all measurements. This is because of the
  # different sampling rates of the data sources in power_sampler
  # and power_metrics.
  means = df.mean(skipna=True)

  # Calculate the standard error of each column.
  stderrs = df.std(skipna=True) / numpy.sqrt(df.count())
  stats = means.to_frame().join(stderrs.to_frame(),
                                lsuffix='mean',
                                rsuffix='stderr')
  stats = stats.rename(columns={"0mean": "mean", "0stderr": "stderr"})

  return stats


def percent_difference(first_value: pd.DataFrame, second_value: pd.DataFrame):
  """
  Returns the comparative percentage difference between two
  values/columns.

  The result is to be read as :
    |second_value| is X% smaller/larger than |first_value|.

  Ex: percent_difference(20, 10) --> -50
  Ex: percent_difference(10, 50) --> 500
  """

  return ((second_value - first_value) / first_value) * 100


def compare(data_dir: str, baseline_summary: str, alternative_summary: str):
  """Open two summary files and compare their values. Saves the results
  in data_dir.

  Args:
    data_dir: The directory to save the comparison csv in.
    baseline_summary: summary.csv for the baseline.
    alternative_summary: summary.csv for the comparison.
  """

  # Get names of the browsers being compared from the paths.
  baseline_name = os.path.basename(
      os.path.dirname(baseline_summary)).split("_")[0]
  alternative_name = os.path.basename(
      os.path.dirname(alternative_summary)).split("_")[0]

  all_stats = []

  # Extract mean and std values for each column of |summary| into a new
  # dataframe.
  baseline_stats = compute_mean_and_stderr(baseline_summary)
  alternative_stats = compute_mean_and_stderr(alternative_summary)

  # Join the calculated values for both browsers into a single dataframe.
  comparison_summary = baseline_stats.join(alternative_stats,
                                           lsuffix=f"_{baseline_name}",
                                           rsuffix=f"_{alternative_name}")

  # Calculate the difference in percent between the baseline and comparison.
  comparison_summary["difference"] = percent_difference(
      baseline_stats["mean"], alternative_stats["mean"])

  # See https://www.cliffsnotes.com/study-guides/statistics/univariate-inferential-tests/two-sample-z-test-for-comparing-two-means
  comparison_summary["z_score"] = (baseline_stats["mean"] -
                                   alternative_stats["mean"]) / numpy.sqrt(
                                       pow(baseline_stats["stderr"], 2) +
                                       pow(alternative_stats["stderr"], 2))

  # See  https://machinelearningmastery.com/critical-values-for-statistical-hypothesis-testing/
  comparison_summary["significance_level"] = scipy_stats.norm.cdf(
      abs(comparison_summary["z_score"]))

  diamond_count = get_diamonds_count(comparison_summary["significance_level"])
  comparison_summary["diamonds"] = diamond_count.apply(get_diamond_string)

  # Drop results for which comparing the mean makes no sense.
  comparison_summary = comparison_summary.drop([
      'battery_max_capacity', 'battery_current_capacity', 'sample_time',
      'elapsed_ns'
  ])

  # Display and save results.
  logging.info(comparison_summary)
  comparison_summary.to_csv(f"{data_dir}/comparison_summary.csv")


def main():
  parser = argparse.ArgumentParser(
      description='Compares two summary files for analysis.')
  parser.add_argument("--output_dir",
                      help="Directory where to write the comparison file.",
                      required=True)
  parser.add_argument("--baseline_dir",
                      help="Directory containing the baseline benchmark data.",
                      required=True)
  parser.add_argument(
      "--alternative_dir",
      help="Directory containing the alternative benchmark data.",
      required=True)
  parser.add_argument('--verbose',
                      action='store_true',
                      help='Print verbose output.')
  args = parser.parse_args()

  if args.verbose:
    log_level = logging.DEBUG
  else:
    log_level = logging.INFO
  logging.basicConfig(format='%(levelname)s: %(message)s', level=log_level)

  baseline_summary_path = os.path.join(args.baseline_dir, "summary.csv")
  alternative_summary_path = os.path.join(args.alternative_dir, "summary.csv")
  summaries = [baseline_summary_path, alternative_summary_path]

  for summary in summaries:
    if not os.path.isfile(summary):
      logging.error(f"summary.csv missing in {summary}.")
      sys.exit(-1)

  compare(args.output_dir, summaries[0], summaries[1])


if __name__ == "__main__":

  # Avoid scientific notation when printing numbers.
  pd.options.display.float_format = '{:.6f}'.format

  main()