#
# 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.
#

"""Unit tests for the transform.util classes."""

from __future__ import absolute_import
from __future__ import division

import logging
import random
import time
import unittest
from builtins import object
from builtins import range

import apache_beam as beam
from apache_beam.coders import coders
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import TestWindowedValue
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import contains_in_any_order
from apache_beam.testing.util import equal_to
from apache_beam.transforms import util
from apache_beam.transforms import window
from apache_beam.transforms.window import GlobalWindow
from apache_beam.transforms.window import GlobalWindows
from apache_beam.transforms.window import IntervalWindow
from apache_beam.transforms.window import Sessions
from apache_beam.transforms.window import SlidingWindows
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils import timestamp
from apache_beam.utils.windowed_value import WindowedValue


class FakeClock(object):

  def __init__(self):
    self._now = time.time()

  def __call__(self):
    return self._now

  def sleep(self, duration):
    self._now += duration


class BatchElementsTest(unittest.TestCase):

  def test_constant_batch(self):
    # Assumes a single bundle...
    with TestPipeline() as p:
      res = (
          p
          | beam.Create(range(35))
          | util.BatchElements(min_batch_size=10, max_batch_size=10)
          | beam.Map(len))
      assert_that(res, equal_to([10, 10, 10, 5]))

  def test_grows_to_max_batch(self):
    # Assumes a single bundle...
    with TestPipeline() as p:
      res = (
          p
          | beam.Create(range(164))
          | util.BatchElements(
              min_batch_size=1, max_batch_size=50, clock=FakeClock())
          | beam.Map(len))
      assert_that(res, equal_to([1, 1, 2, 4, 8, 16, 32, 50, 50]))

  def test_windowed_batches(self):
    # Assumes a single bundle, in order...
    with TestPipeline() as p:
      res = (
          p
          | beam.Create(range(47))
          | beam.Map(lambda t: window.TimestampedValue(t, t))
          | beam.WindowInto(window.FixedWindows(30))
          | util.BatchElements(
              min_batch_size=5, max_batch_size=10, clock=FakeClock())
          | beam.Map(len))
      assert_that(res, equal_to([
          5, 5, 10, 10,  # elements in [0, 30)
          10, 7,         # elements in [30, 47)
      ]))

  def test_target_duration(self):
    clock = FakeClock()
    batch_estimator = util._BatchSizeEstimator(
        target_batch_overhead=None, target_batch_duration_secs=10, clock=clock)
    batch_duration = lambda batch_size: 1 + .7 * batch_size
    # 1 + 12 * .7 is as close as we can get to 10 as possible.
    expected_sizes = [1, 2, 4, 8, 12, 12, 12]
    actual_sizes = []
    for _ in range(len(expected_sizes)):
      actual_sizes.append(batch_estimator.next_batch_size())
      with batch_estimator.record_time(actual_sizes[-1]):
        clock.sleep(batch_duration(actual_sizes[-1]))
    self.assertEqual(expected_sizes, actual_sizes)

  def test_target_overhead(self):
    clock = FakeClock()
    batch_estimator = util._BatchSizeEstimator(
        target_batch_overhead=.05, target_batch_duration_secs=None, clock=clock)
    batch_duration = lambda batch_size: 1 + .7 * batch_size
    # At 27 items, a batch takes ~20 seconds with 5% (~1 second) overhead.
    expected_sizes = [1, 2, 4, 8, 16, 27, 27, 27]
    actual_sizes = []
    for _ in range(len(expected_sizes)):
      actual_sizes.append(batch_estimator.next_batch_size())
      with batch_estimator.record_time(actual_sizes[-1]):
        clock.sleep(batch_duration(actual_sizes[-1]))
    self.assertEqual(expected_sizes, actual_sizes)

  def test_variance(self):
    clock = FakeClock()
    variance = 0.25
    batch_estimator = util._BatchSizeEstimator(
        target_batch_overhead=.05, target_batch_duration_secs=None,
        variance=variance, clock=clock)
    batch_duration = lambda batch_size: 1 + .7 * batch_size
    expected_target = 27
    actual_sizes = []
    for _ in range(util._BatchSizeEstimator._MAX_DATA_POINTS - 1):
      actual_sizes.append(batch_estimator.next_batch_size())
      with batch_estimator.record_time(actual_sizes[-1]):
        clock.sleep(batch_duration(actual_sizes[-1]))
    # Check that we're testing a good range of values.
    stable_set = set(actual_sizes[-20:])
    self.assertGreater(len(stable_set), 3)
    self.assertGreater(
        min(stable_set), expected_target - expected_target * variance)
    self.assertLess(
        max(stable_set), expected_target + expected_target * variance)

  def _run_regression_test(self, linear_regression_fn, test_outliers):
    xs = [random.random() for _ in range(10)]
    ys = [2*x + 1 for x in xs]
    a, b = linear_regression_fn(xs, ys)
    self.assertAlmostEqual(a, 1)
    self.assertAlmostEqual(b, 2)

    xs = [1 + random.random() for _ in range(100)]
    ys = [7*x + 5 + 0.01 * random.random() for x in xs]
    a, b = linear_regression_fn(xs, ys)
    self.assertAlmostEqual(a, 5, delta=0.02)
    self.assertAlmostEqual(b, 7, delta=0.02)

    # Test repeated xs
    xs = [1 + random.random()] * 100
    ys = [7 * x + 5 + 0.01 * random.random() for x in xs]
    a, b = linear_regression_fn(xs, ys)
    self.assertAlmostEqual(a, 0, delta=0.02)
    self.assertAlmostEqual(
        b, sum(ys)/(len(ys) * xs[0]), delta=0.02)

    if test_outliers:
      xs = [1 + random.random() for _ in range(100)]
      ys = [2*x + 1 for x in xs]
      a, b = linear_regression_fn(xs, ys)
      self.assertAlmostEqual(a, 1)
      self.assertAlmostEqual(b, 2)

      # An outlier or two doesn't affect the result.
      for _ in range(2):
        xs += [10]
        ys += [30]
        a, b = linear_regression_fn(xs, ys)
        self.assertAlmostEqual(a, 1)
        self.assertAlmostEqual(b, 2)

      # But enough of them, and they're no longer outliers.
      xs += [10] * 10
      ys += [30] * 10
      a, b = linear_regression_fn(xs, ys)
      self.assertLess(a, 0.5)
      self.assertGreater(b, 2.5)

  def test_no_numpy_regression(self):
    self._run_regression_test(
        util._BatchSizeEstimator.linear_regression_no_numpy, False)

  def test_numpy_regression(self):
    try:
      # pylint: disable=wrong-import-order, wrong-import-position
      import numpy as _
    except ImportError:
      self.skipTest('numpy not available')
    self._run_regression_test(
        util._BatchSizeEstimator.linear_regression_numpy, True)


class IdentityWindowTest(unittest.TestCase):

  def test_window_preserved(self):
    expected_timestamp = timestamp.Timestamp(5)
    expected_window = window.IntervalWindow(1.0, 2.0)

    class AddWindowDoFn(beam.DoFn):
      def process(self, element):
        yield WindowedValue(
            element, expected_timestamp, [expected_window])

    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
    expected_windows = [
        TestWindowedValue(kv, expected_timestamp, [expected_window])
        for kv in data]
    before_identity = (pipeline
                       | 'start' >> beam.Create(data)
                       | 'add_windows' >> beam.ParDo(AddWindowDoFn()))
    assert_that(before_identity, equal_to(expected_windows),
                label='before_identity', reify_windows=True)
    after_identity = (before_identity
                      | 'window' >> beam.WindowInto(
                          beam.transforms.util._IdentityWindowFn(
                              coders.IntervalWindowCoder())))
    assert_that(after_identity, equal_to(expected_windows),
                label='after_identity', reify_windows=True)
    pipeline.run()

  def test_no_window_context_fails(self):
    expected_timestamp = timestamp.Timestamp(5)
    # Assuming the default window function is window.GlobalWindows.
    expected_window = window.GlobalWindow()

    class AddTimestampDoFn(beam.DoFn):
      def process(self, element):
        yield window.TimestampedValue(element, expected_timestamp)

    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
    expected_windows = [
        TestWindowedValue(kv, expected_timestamp, [expected_window])
        for kv in data]
    before_identity = (pipeline
                       | 'start' >> beam.Create(data)
                       | 'add_timestamps' >> beam.ParDo(AddTimestampDoFn()))
    assert_that(before_identity, equal_to(expected_windows),
                label='before_identity', reify_windows=True)
    after_identity = (before_identity
                      | 'window' >> beam.WindowInto(
                          beam.transforms.util._IdentityWindowFn(
                              coders.GlobalWindowCoder()))
                      # This DoFn will return TimestampedValues, making
                      # WindowFn.AssignContext passed to IdentityWindowFn
                      # contain a window of None. IdentityWindowFn should
                      # raise an exception.
                      | 'add_timestamps2' >> beam.ParDo(AddTimestampDoFn()))
    assert_that(after_identity, equal_to(expected_windows),
                label='after_identity', reify_windows=True)
    with self.assertRaisesRegexp(ValueError, r'window.*None.*add_timestamps2'):
      pipeline.run()


class ReshuffleTest(unittest.TestCase):

  def test_reshuffle_contents_unchanged(self):
    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]
    result = (pipeline
              | beam.Create(data)
              | beam.Reshuffle())
    assert_that(result, equal_to(data))
    pipeline.run()

  def test_reshuffle_after_gbk_contents_unchanged(self):
    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]
    expected_result = [(1, [1, 2, 3]), (2, [1, 2]), (3, [1])]

    after_gbk = (pipeline
                 | beam.Create(data)
                 | beam.GroupByKey())
    assert_that(after_gbk, equal_to(expected_result), label='after_gbk')
    after_reshuffle = after_gbk | beam.Reshuffle()
    assert_that(after_reshuffle, equal_to(expected_result),
                label='after_reshuffle')
    pipeline.run()

  def test_reshuffle_timestamps_unchanged(self):
    pipeline = TestPipeline()
    timestamp = 5
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]
    expected_result = [TestWindowedValue(v, timestamp, [GlobalWindow()])
                       for v in data]
    before_reshuffle = (pipeline
                        | 'start' >> beam.Create(data)
                        | 'add_timestamp' >> beam.Map(
                            lambda v: beam.window.TimestampedValue(v,
                                                                   timestamp)))
    assert_that(before_reshuffle, equal_to(expected_result),
                label='before_reshuffle', reify_windows=True)
    after_reshuffle = before_reshuffle | beam.Reshuffle()
    assert_that(after_reshuffle, equal_to(expected_result),
                label='after_reshuffle', reify_windows=True)
    pipeline.run()

  def test_reshuffle_windows_unchanged(self):
    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
    expected_data = [TestWindowedValue(v, t, [w]) for (v, t, w) in [
        ((1, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
        ((2, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
        ((3, [1]), 3.0, IntervalWindow(1.0, 3.0)),
        ((1, [4]), 6.0, IntervalWindow(4.0, 6.0))]]
    before_reshuffle = (pipeline
                        | 'start' >> beam.Create(data)
                        | 'add_timestamp' >> beam.Map(
                            lambda v: beam.window.TimestampedValue(v, v[1]))
                        | 'window' >> beam.WindowInto(Sessions(gap_size=2))
                        | 'group_by_key' >> beam.GroupByKey())
    assert_that(before_reshuffle, equal_to(expected_data),
                label='before_reshuffle', reify_windows=True)
    after_reshuffle = before_reshuffle | beam.Reshuffle()
    assert_that(after_reshuffle, equal_to(expected_data),
                label='after reshuffle', reify_windows=True)
    pipeline.run()

  def test_reshuffle_window_fn_preserved(self):
    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
    expected_windows = [TestWindowedValue(v, t, [w]) for (v, t, w) in [
        ((1, 1), 1.0, IntervalWindow(1.0, 3.0)),
        ((2, 1), 1.0, IntervalWindow(1.0, 3.0)),
        ((3, 1), 1.0, IntervalWindow(1.0, 3.0)),
        ((1, 2), 2.0, IntervalWindow(2.0, 4.0)),
        ((2, 2), 2.0, IntervalWindow(2.0, 4.0)),
        ((1, 4), 4.0, IntervalWindow(4.0, 6.0))]]
    expected_merged_windows = [TestWindowedValue(v, t, [w]) for (v, t, w) in [
        ((1, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
        ((2, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
        ((3, [1]), 3.0, IntervalWindow(1.0, 3.0)),
        ((1, [4]), 6.0, IntervalWindow(4.0, 6.0))]]
    before_reshuffle = (pipeline
                        | 'start' >> beam.Create(data)
                        | 'add_timestamp' >> beam.Map(
                            lambda v: TimestampedValue(v, v[1]))
                        | 'window' >> beam.WindowInto(Sessions(gap_size=2)))
    assert_that(before_reshuffle, equal_to(expected_windows),
                label='before_reshuffle', reify_windows=True)
    after_reshuffle = before_reshuffle | beam.Reshuffle()
    assert_that(after_reshuffle, equal_to(expected_windows),
                label='after_reshuffle', reify_windows=True)
    after_group = after_reshuffle | beam.GroupByKey()
    assert_that(after_group, equal_to(expected_merged_windows),
                label='after_group', reify_windows=True)
    pipeline.run()

  def test_reshuffle_global_window(self):
    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
    expected_data = [(1, [1, 2, 4]), (2, [1, 2]), (3, [1])]
    before_reshuffle = (pipeline
                        | beam.Create(data)
                        | beam.WindowInto(GlobalWindows())
                        | beam.GroupByKey())
    assert_that(before_reshuffle, equal_to(expected_data),
                label='before_reshuffle')
    after_reshuffle = before_reshuffle | beam.Reshuffle()
    assert_that(after_reshuffle, equal_to(expected_data),
                label='after reshuffle')
    pipeline.run()

  def test_reshuffle_sliding_window(self):
    pipeline = TestPipeline()
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
    window_size = 2
    expected_data = [(1, [1, 2, 4]), (2, [1, 2]), (3, [1])] * window_size
    before_reshuffle = (pipeline
                        | beam.Create(data)
                        | beam.WindowInto(SlidingWindows(
                            size=window_size, period=1))
                        | beam.GroupByKey())
    assert_that(before_reshuffle, equal_to(expected_data),
                label='before_reshuffle')
    after_reshuffle = before_reshuffle | beam.Reshuffle()
    # If Reshuffle applies the sliding window function a second time there
    # should be extra values for each key.
    assert_that(after_reshuffle, equal_to(expected_data),
                label='after reshuffle')
    pipeline.run()

  def test_reshuffle_streaming_global_window(self):
    options = PipelineOptions()
    options.view_as(StandardOptions).streaming = True
    pipeline = TestPipeline(options=options)
    data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
    expected_data = [(1, [1, 2, 4]), (2, [1, 2]), (3, [1])]
    before_reshuffle = (pipeline
                        | beam.Create(data)
                        | beam.WindowInto(GlobalWindows())
                        | beam.GroupByKey())
    assert_that(before_reshuffle, equal_to(expected_data),
                label='before_reshuffle')
    after_reshuffle = before_reshuffle | beam.Reshuffle()
    assert_that(after_reshuffle, equal_to(expected_data),
                label='after reshuffle')
    pipeline.run()


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()