#!/usr/bin/env python3

# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""This script is used to run the error_calculator binary.

  ./error_calculator.py 1 657579885 14 1

This script assumes it is being run from a third_party embedding in a fuchsia
checkout. If you are running from the cobalt standalone repository, run
'cobaltb.py calculate_error' instead. If the build configuration doesn't match
the assumed fuchsia or cobalt checkout, use the --bin_dir flag to specify the
path to the error_calculator and config_parser binaries.

A new version of the registry is generated every time.
"""

import subprocess
import os
import argparse
import sys

THIS_DIR = os.path.abspath(os.path.dirname(__file__))
COBALT_ROOT_DIR = os.path.abspath(os.path.join(THIS_DIR, '..'))
PRIVACY_PARAMS = os.path.join(COBALT_ROOT_DIR, 'src', 'algorithms', 'privacy',
                              'data', 'privacy_encoding_params')
REGISTRY_PROTO_DEFAULT = os.path.join(COBALT_ROOT_DIR, 'out', 'registry.pb')
POPULATION_DEFAULT = 100000


def estimate_from_args(error_binary, args):
  """Unpacks and validates command line arguments before calling error calculator.

  Args:
    error_binary: The filepath of the error calculator binary.
    args: Python arguments defined below.
  """
  _validate_args(args)
  estimate(error_binary, args.registry_proto, args.customer_id, args.project_id,
           args.metric_id, args.report_id, args.population, args.epsilon)


def estimate(error_binary,
             registry_proto,
             customer_id,
             project_id,
             metric_id,
             report_id,
             population,
             epsilon=None):
  """Calls the error_calculator binary with the specified arguments.

  Args:
    error_binary: The string filepath to the error calculator binary.
    registry_proto: The string filepath to the registry formatted as a
      serialized proto.
    customer_id: The integer customer id.
    project_id: The integer project id.
    metric_id: The integer metric id.
    report_id: The integer report id.
    population: Integer value estimating the number of reporting devices.
    epsilon: Optional; Epsilon value for which to estimate error.
  """
  error_args = [
      '-registry_proto', registry_proto, '--privacy_params', PRIVACY_PARAMS,
      '--population',
      str(population)
  ]
  if epsilon:
    error_args = error_args + ['--epsilon', str(epsilon)]
  error_args = error_args + [
      str(customer_id),
      str(project_id),
      str(metric_id),
      str(report_id)
  ]
  subprocess.check_call([error_binary] + error_args)


# TODO(jaredweinstein): Use the build system to generate the registry only when
# necessary.
def generate_registry(registry_proto, config_dir, config_parser):
  """Generates a binary encoding of the cobalt registry.

  Args:
    registry_proto: The filepath of the error calculator binary. This is the
      output of the config_parser and the input for the error_calculator
    config_dir: Location of the cobalt config.
    config_parser: Location of the binary used to generate the registry_proto.
  """
  if not config_parser:
    sys.exit('Run \'config_parser --out_filename %s\' and try again.' %
             registry_proto)
  subprocess.check_call([
      config_parser, '--output_file', registry_proto, '--config_dir',
      config_dir, '--privacy_encoding_params_file', PRIVACY_PARAMS
  ])
  print('Wrote binary encoding of registry to %s.\n' % registry_proto)


def add_parse_args(parser):
  """Adds the standard arguments required for the error_calculator.

  This is used to set arguments for this script and the cobaltb.py script.

  Args:
    parser: An ArgumentParser to be augmented with error calculator arguments.
  """
  parser.add_argument(
      '--registry_proto',
      help='Set a specific filepath for the binary encoding of the Cobalt Registry. Default: %s'
      % REGISTRY_PROTO_DEFAULT,
      default=REGISTRY_PROTO_DEFAULT)
  parser.add_argument(
      '--population',
      help='Expected number of devices contributing to the report. Default: %s'
      % POPULATION_DEFAULT,
      default=POPULATION_DEFAULT)
  parser.add_argument(
      '--bin_dir',
      help='Directory containing the error_calculator and config_parser binaries. '
      'If unset, the script attempts to find binaries based on the default '
      'build configuration for Fuchsia and Cobalt.',
      default=None)
  parser.add_argument(
      '--epsilon',
      help='If set, estimates the error using the specified epsilon value.')
  parser.add_argument(
      'customer_id', help='a report\'s parent customer id', type=int)
  parser.add_argument(
      'project_id', help='a report\'s parent project id.', type=int)
  parser.add_argument(
      'metric_id', help='a report\'s parent metric id.', type=int)
  parser.add_argument(
      'report_id', help='a report\'s parent report id.', type=int)


def _validate_args(args):
  if args.registry_proto == None:
    sys.exit('--registry_proto is required')
  if args.population == None:
    sys.exit('--population flag is required')
  if not os.path.exists(args.registry_proto):
    sys.exit(
        'No serialized proto found. Try running \'config_parser --out_filename %s\' and try again.'
        % args.registry_proto)


if __name__ == '__main__':
  # Assumes the user is running from a Fuchsia checkout.
  # If you're running this script directly from a cobalt checkout, use
  # `cobaltb.py calculate_error` instead.
  parser = argparse.ArgumentParser()
  add_parse_args(parser)
  args = parser.parse_args()

  if args.bin_dir:
    bin_dir = args.bin_dir
  else:
    out = subprocess.check_output(['fx', 'get-build-dir'])
    out_dir = out.decode('utf-8').rstrip()
    bin_dir = os.path.join(out_dir, 'host_x64')

  error_binary = os.path.join(bin_dir, 'error_calculator')
  config_parser = os.path.join(bin_dir, 'config_parser')
  config_dir = os.path.join(COBALT_ROOT_DIR, '..', 'cobalt_config')
  config_dir = os.path.abspath(config_dir)
  if not os.path.exists(error_binary):
    sys.exit(
        'Error Calculator binary not found: %s.\nRun \'fx build\' and try again.'
        % error_binary)

  generate_registry(args.registry_proto, config_dir, config_parser)
  estimate_from_args(error_binary, args)