# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""A microbenchmark for measuring performance of coders.

This runs a sequence of encode-decode operations on random inputs
to collect performance of various coders.

To evaluate coders performance we approximate the behavior
how the coders are used in PCollections: we encode and decode
a list of elements. An element can be a string, a list of integers,
a windowed value, or any other object we want a coder to process.

Run as:
  python -m apache_beam.tools.coders_microbenchmark

"""

# pytype: skip-file

import argparse
import logging
import random
import re
import string
import sys

import apache_beam as beam
from apache_beam.coders import coder_impl
from apache_beam.coders import coders
from apache_beam.coders import coders_test_common
from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message
from apache_beam.coders import row_coder
from apache_beam.coders import typecoders
from apache_beam.tools import utils
from apache_beam.transforms import window
from apache_beam.typehints import trivial_inference
from apache_beam.typehints.pandas_type_compatibility import DataFrameBatchConverterDropIndex
from apache_beam.utils import windowed_value


def coder_benchmark_factory(coder, generate_fn):
  """Creates a benchmark that encodes and decodes a list of elements.

  Args:
    coder: coder to use to encode an element.
    generate_fn: a callable that generates an element.
  """
  class CoderBenchmark(object):
    def __init__(self, num_elements_per_benchmark):
      self._coder = coders.IterableCoder(coder)
      self._list = [generate_fn() for _ in range(num_elements_per_benchmark)]

    def __call__(self):
      # Calling coder operations on a single element at a time may incur
      # unrelevant overhead. To compensate, we use a list elements.
      _ = self._coder.decode(self._coder.encode(self._list))

  CoderBenchmark.__name__ = "%s, %s" % (generate_fn.__name__, str(coder))

  return CoderBenchmark


def batch_row_coder_benchmark_factory(generate_fn, use_batch):
  """Creates a benchmark that encodes and decodes a list of elements.

  Args:
    coder: coder to use to encode an element.
    generate_fn: a callable that generates an element.
  """
  class CoderBenchmark(object):
    def __init__(self, num_elements_per_benchmark):
      self._use_batch = use_batch
      row_instance = generate_fn()
      row_type = trivial_inference.instance_to_type(row_instance)
      self._row_coder = get_row_coder(row_instance)
      self._batch_converter = DataFrameBatchConverterDropIndex(row_type)
      self._seq_coder = coders.IterableCoder(self._row_coder)
      self._data = self._batch_converter.produce_batch(
          [generate_fn() for _ in range(num_elements_per_benchmark)])

    def __call__(self):
      if self._use_batch:
        impl = self._row_coder.get_impl()
        columnar = {
            col: self._data[col].to_numpy()
            for col in self._data.columns
        }
        output_stream = coder_impl.create_OutputStream()
        impl.encode_batch_to_stream(columnar, output_stream)
        impl.decode_batch_from_stream(
            columnar, coder_impl.create_InputStream(output_stream.get()))

      else:
        # Calling coder operations on a single element at a time may incur
        # unrelevant overhead. To compensate, we use a list elements.
        self._batch_converter.produce_batch(
            self._seq_coder.decode(
                self._seq_coder.encode(
                    self._batch_converter.explode_batch(self._data))))

  CoderBenchmark.__name__ = "%s, BatchRowCoder%s" % (
      generate_fn.__name__, use_batch)

  return CoderBenchmark


def small_int():
  return random.randint(0, 127)


def large_int():
  return random.randint(sys.maxsize >> 2, sys.maxsize)


def random_string(length):
  return ''.join(
      random.choice(string.ascii_letters + string.digits)
      for _ in range(length))


def small_string():
  return random_string(4)


def large_string():
  return random_string(100)


def list_int(size):
  return [small_int() for _ in range(size)]


def dict_int_int(size):
  return {i: i for i in list_int(size)}


def small_list():
  return list_int(10)


def large_list():
  # Bool is the last item in FastPrimitiveCoders before pickle.
  return [bool(k) for k in list_int(1000)]


def small_tuple():
  # Benchmark a common case of 2-element tuples.
  return tuple(list_int(2))


def large_tuple():
  return tuple(large_list())


def small_dict():
  return {i: i for i in small_list()}


def large_dict():
  return {i: i for i in large_list()}


def large_iterable():
  yield 'a' * coders.coder_impl.SequenceCoderImpl._DEFAULT_BUFFER_SIZE
  for k in range(1000):
    yield k


def random_message_with_map(size):
  message = test_message.MessageWithMap()
  keys = list_int(size)
  random.shuffle(keys)
  for key in keys:
    message.field1[str(key)].field1 = small_string()
  return message


def small_message_with_map():
  return random_message_with_map(5)


def large_message_with_map():
  return random_message_with_map(20)


def globally_windowed_value():
  return windowed_value.WindowedValue(
      value=small_int(), timestamp=12345678, windows=(window.GlobalWindow(), ))


def random_windowed_value(num_windows):
  return windowed_value.WindowedValue(
      value=small_int(),
      timestamp=12345678,
      windows=tuple(
          window.IntervalWindow(i * 10, i * 10 + small_int())
          for i in range(num_windows)))


def wv_with_one_window():
  return random_windowed_value(num_windows=1)


def wv_with_multiple_windows():
  return random_windowed_value(num_windows=32)


def tiny_row():
  return beam.Row(int_value=1)


def large_row():
  return beam.Row(**{f'int_{ix}': ix for ix in range(20)})


def nullable_row():
  return beam.Row(**{f'int_{ix}': ix if ix % 2 else None for ix in range(20)})


def diverse_row():
  return beam.Row(
      int_value=1,
      float_value=3.14159,
      str_value='beam',
      row_value=beam.Row(int_value=2, float_value=2.718281828))


def get_row_coder(row_instance):
  coder = typecoders.registry.get_coder(
      trivial_inference.instance_to_type(row_instance))
  assert isinstance(coder, row_coder.RowCoder)
  return coder


def row_coder_benchmark_factory(generate_fn):
  return coder_benchmark_factory(get_row_coder(generate_fn()), generate_fn)


def importable_named_tuple():
  return [coders_test_common.MyTypedNamedTuple('a', i) for i in range(1000)]


def run_coder_benchmarks(
    num_runs, input_size, seed, verbose, filter_regex='.*'):
  random.seed(seed)

  # TODO(https://github.com/apache/beam/issues/18788): Pick coders using type
  # hints, for example:
  # tuple_coder = typecoders.registry.get_coder(typing.Tuple[int, ...])
  benchmarks = [
      coder_benchmark_factory(coders.FastPrimitivesCoder(), small_int),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), large_int),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), small_string),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), large_string),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), small_list),
      coder_benchmark_factory(
          coders.IterableCoder(coders.FastPrimitivesCoder()), small_list),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), large_list),
      coder_benchmark_factory(
          coders.IterableCoder(coders.FastPrimitivesCoder()), large_list),
      coder_benchmark_factory(
          coders.IterableCoder(coders.FastPrimitivesCoder()), large_iterable),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), small_tuple),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), large_tuple),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), small_dict),
      coder_benchmark_factory(coders.FastPrimitivesCoder(), large_dict),
      coder_benchmark_factory(
          coders.ProtoCoder(test_message.MessageWithMap),
          small_message_with_map),
      coder_benchmark_factory(
          coders.ProtoCoder(test_message.MessageWithMap),
          large_message_with_map),
      coder_benchmark_factory(
          coders.DeterministicProtoCoder(test_message.MessageWithMap),
          small_message_with_map),
      coder_benchmark_factory(
          coders.DeterministicProtoCoder(test_message.MessageWithMap),
          large_message_with_map),
      coder_benchmark_factory(
          coders.WindowedValueCoder(coders.FastPrimitivesCoder()),
          wv_with_one_window),
      coder_benchmark_factory(
          coders.WindowedValueCoder(
              coders.FastPrimitivesCoder(), coders.IntervalWindowCoder()),
          wv_with_multiple_windows),
      coder_benchmark_factory(
          coders.WindowedValueCoder(
              coders.FastPrimitivesCoder(), coders.GlobalWindowCoder()),
          globally_windowed_value),
      coder_benchmark_factory(
          coders.LengthPrefixCoder(coders.FastPrimitivesCoder()), small_int),
      row_coder_benchmark_factory(tiny_row),
      row_coder_benchmark_factory(large_row),
      row_coder_benchmark_factory(nullable_row),
      row_coder_benchmark_factory(diverse_row),
      batch_row_coder_benchmark_factory(tiny_row, False),
      batch_row_coder_benchmark_factory(tiny_row, True),
      batch_row_coder_benchmark_factory(large_row, False),
      batch_row_coder_benchmark_factory(large_row, True),
      batch_row_coder_benchmark_factory(nullable_row, False),
      batch_row_coder_benchmark_factory(nullable_row, True),
      batch_row_coder_benchmark_factory(diverse_row, False),
      batch_row_coder_benchmark_factory(diverse_row, True),
      coder_benchmark_factory(
          coders.IterableCoder(
              coders.FastPrimitivesCoder().as_deterministic_coder(
                  step_label="step")),
          importable_named_tuple),
  ]

  suite = [
      utils.BenchmarkConfig(b, input_size, num_runs) for b in benchmarks
      if re.search(filter_regex, b.__name__, flags=re.I)
  ]
  utils.run_benchmarks(suite, verbose=verbose)


if __name__ == "__main__":
  logging.basicConfig()

  parser = argparse.ArgumentParser()
  parser.add_argument('--filter', default='.*')
  parser.add_argument('--num_runs', default=20, type=int)
  parser.add_argument('--num_elements_per_benchmark', default=1000, type=int)
  parser.add_argument('--seed', default=42, type=int)
  options = parser.parse_args()

  utils.check_compiled("apache_beam.coders.coder_impl")

  num_runs = 20
  num_elements_per_benchmark = 1000
  seed = 42  # Fix the seed for better consistency

  run_coder_benchmarks(
      options.num_runs,
      options.num_elements_per_benchmark,
      options.seed,
      verbose=True,
      filter_regex=options.filter)