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

# pytype: skip-file

import copy
import itertools
import random
import threading
import unittest

from apache_beam.metrics.cells import BoundedTrieData
from apache_beam.metrics.cells import CounterCell
from apache_beam.metrics.cells import DistributionCell
from apache_beam.metrics.cells import DistributionData
from apache_beam.metrics.cells import GaugeCell
from apache_beam.metrics.cells import GaugeData
from apache_beam.metrics.cells import HistogramCell
from apache_beam.metrics.cells import HistogramCellFactory
from apache_beam.metrics.cells import HistogramData
from apache_beam.metrics.cells import StringSetCell
from apache_beam.metrics.cells import StringSetData
from apache_beam.metrics.cells import _BoundedTrieNode
from apache_beam.metrics.metricbase import MetricName
from apache_beam.utils.histogram import Histogram
from apache_beam.utils.histogram import LinearBucket


class TestCounterCell(unittest.TestCase):
  @classmethod
  def _modify_counter(cls, d):
    for i in range(cls.NUM_ITERATIONS):
      d.inc(i)

  NUM_THREADS = 5
  NUM_ITERATIONS = 100

  def test_parallel_access(self):
    # We create NUM_THREADS threads that concurrently modify the counter.
    threads = []
    c = CounterCell()
    for _ in range(TestCounterCell.NUM_THREADS):
      t = threading.Thread(
          target=TestCounterCell._modify_counter, args=(c, ))
      threads.append(t)
      t.start()

    for t in threads:
      t.join()

    total = (
        self.NUM_ITERATIONS * (self.NUM_ITERATIONS - 1) // 2 * self.NUM_THREADS)
    self.assertEqual(c.get_cumulative(), total)

  def test_basic_operations(self):
    c = CounterCell()
    c.inc(2)
    self.assertEqual(c.get_cumulative(), 2)

    c.dec(10)
    self.assertEqual(c.get_cumulative(), -8)

    c.dec()
    self.assertEqual(c.get_cumulative(), -9)

    c.inc()
    self.assertEqual(c.get_cumulative(), -8)

  def test_start_time_set(self):
    c = CounterCell()
    c.inc(2)

    name = MetricName('namespace', 'name1')
    mi = c.to_runner_api_monitoring_info(name, 'transform_id')
    self.assertGreater(mi.start_time.seconds, 0)


class TestDistributionCell(unittest.TestCase):
  @classmethod
  def _modify_distribution(cls, d):
    for i in range(cls.NUM_ITERATIONS):
      d.update(i)

  NUM_THREADS = 5
  NUM_ITERATIONS = 100

  def test_parallel_access(self):
    # We create NUM_THREADS threads that concurrently modify the distribution.
    threads = []
    d = DistributionCell()
    for _ in range(TestDistributionCell.NUM_THREADS):
      t = threading.Thread(
          target=TestDistributionCell._modify_distribution, args=(d, ))
      threads.append(t)
      t.start()

    for t in threads:
      t.join()

    total = (
        self.NUM_ITERATIONS * (self.NUM_ITERATIONS - 1) // 2 * self.NUM_THREADS)

    count = (self.NUM_ITERATIONS * self.NUM_THREADS)

    self.assertEqual(
        d.get_cumulative(),
        DistributionData(total, count, 0, self.NUM_ITERATIONS - 1))

  def test_basic_operations(self):
    d = DistributionCell()
    d.update(10)
    self.assertEqual(d.get_cumulative(), DistributionData(10, 1, 10, 10))

    d.update(2)
    self.assertEqual(d.get_cumulative(), DistributionData(12, 2, 2, 10))

    d.update(900)
    self.assertEqual(d.get_cumulative(), DistributionData(912, 3, 2, 900))

  def test_integer_only(self):
    d = DistributionCell()
    d.update(3.1)
    d.update(3.2)
    d.update(3.3)
    self.assertEqual(d.get_cumulative(), DistributionData(9, 3, 3, 3))

  def test_start_time_set(self):
    d = DistributionCell()
    d.update(3.1)

    name = MetricName('namespace', 'name1')
    mi = d.to_runner_api_monitoring_info(name, 'transform_id')
    self.assertGreater(mi.start_time.seconds, 0)


class TestGaugeCell(unittest.TestCase):
  def test_basic_operations(self):
    g = GaugeCell()
    g.set(10)
    self.assertEqual(g.get_cumulative().value, GaugeData(10).value)

    g.set(2)
    self.assertEqual(g.get_cumulative().value, 2)

  def test_integer_only(self):
    g = GaugeCell()
    g.set(3.3)
    self.assertEqual(g.get_cumulative().value, 3)

  def test_combine_appropriately(self):
    g1 = GaugeCell()
    g1.set(3)

    g2 = GaugeCell()
    g2.set(1)

    # THe second Gauge, with value 1 was the most recent, so it should be
    # the final result.
    result = g2.combine(g1)
    self.assertEqual(result.data.value, 1)

  def test_start_time_set(self):
    g1 = GaugeCell()
    g1.set(3)

    name = MetricName('namespace', 'name1')
    mi = g1.to_runner_api_monitoring_info(name, 'transform_id')
    self.assertGreater(mi.start_time.seconds, 0)


class TestStringSetCell(unittest.TestCase):
  def test_not_leak_mutable_set(self):
    c = StringSetCell()
    c.add('test')
    c.add('another')
    s = c.get_cumulative()
    self.assertEqual(s, StringSetData({'test', 'another'}, 11))
    s.add('yet another')
    self.assertEqual(c.get_cumulative(), StringSetData({'test', 'another'}, 11))

  def test_combine_appropriately(self):
    s1 = StringSetCell()
    s1.add('1')
    s1.add('2')

    s2 = StringSetCell()
    s2.add('1')
    s2.add('3')

    result = s2.combine(s1)
    self.assertEqual(result.data, StringSetData({'1', '2', '3'}))

  def test_add_size_tracked_correctly(self):
    s = StringSetCell()
    s.add('1')
    s.add('2')
    self.assertEqual(s.data.string_size, 2)
    s.add('2')
    s.add('3')
    self.assertEqual(s.data.string_size, 3)


class TestBoundedTrieNode(unittest.TestCase):
  @classmethod
  def random_segments_fixed_depth(cls, n, depth, overlap, rand):
    if depth == 0:
      yield from ((), ) * n
    else:
      seen = []
      to_string = lambda ix: chr(ord('a') + ix) if ix < 26 else f'z{ix}'
      for suffix in cls.random_segments_fixed_depth(n, depth - 1, overlap,
                                                    rand):
        if not seen or rand.random() > overlap:
          prefix = to_string(len(seen))
          seen.append(prefix)
        else:
          prefix = rand.choice(seen)
        yield (prefix, ) + suffix

  @classmethod
  def random_segments(cls, n, min_depth, max_depth, overlap, rand):
    for depth, segments in zip(
        itertools.cycle(range(min_depth, max_depth + 1)),
        cls.random_segments_fixed_depth(n, max_depth, overlap, rand)):
      yield segments[:depth]

  def assert_covers(self, node, expected, max_truncated=0):
    self.assert_covers_flattened(node.flattened(), expected, max_truncated)

  def assert_covers_flattened(self, flattened, expected, max_truncated=0):
    expected = set(expected)
    # Split node into the exact and truncated segments.
    partitioned = {True: set(), False: set()}
    for segments in flattened:
      partitioned[segments[-1]].add(segments[:-1])
    exact, truncated = partitioned[False], partitioned[True]
    # Check we cover both parts.
    self.assertLessEqual(len(truncated), max_truncated, truncated)
    self.assertTrue(exact.issubset(expected), exact - expected)
    seen_truncated = set()
    for segments in expected - exact:
      found = 0
      for ix in range(len(segments)):
        if segments[:ix] in truncated:
          seen_truncated.add(segments[:ix])
          found += 1
      if found != 1:
        self.fail(
            f"Expected exactly one prefix of {segments} "
            f"to occur in {truncated}, found {found}")
    self.assertEqual(seen_truncated, truncated, truncated - seen_truncated)

  def run_covers_test(self, flattened, expected, max_truncated):
    def parse(s):
      return tuple(s.strip('*')) + (s.endswith('*'), )

    self.assert_covers_flattened([parse(s) for s in flattened],
                                 [tuple(s) for s in expected],
                                 max_truncated)

  def test_covers_exact(self):
    self.run_covers_test(['ab', 'ac', 'cd'], ['ab', 'ac', 'cd'], 0)
    with self.assertRaises(AssertionError):
      self.run_covers_test(['ab', 'ac', 'cd'], ['ac', 'cd'], 0)
    with self.assertRaises(AssertionError):
      self.run_covers_test(['ab', 'ac'], ['ab', 'ac', 'cd'], 0)
    with self.assertRaises(AssertionError):
      self.run_covers_test(['a*', 'cd'], ['ab', 'ac', 'cd'], 0)

  def test_covers_trunacted(self):
    self.run_covers_test(['a*', 'cd'], ['ab', 'ac', 'cd'], 1)
    self.run_covers_test(['a*', 'cd'], ['ab', 'ac', 'abcde', 'cd'], 1)
    with self.assertRaises(AssertionError):
      self.run_covers_test(['ab', 'ac', 'cd'], ['ac', 'cd'], 1)
    with self.assertRaises(AssertionError):
      self.run_covers_test(['ab', 'ac'], ['ab', 'ac', 'cd'], 1)
    with self.assertRaises(AssertionError):
      self.run_covers_test(['a*', 'c*'], ['ab', 'ac', 'cd'], 1)
    with self.assertRaises(AssertionError):
      self.run_covers_test(['a*', 'c*'], ['ab', 'ac'], 1)

  def run_test(self, to_add):
    everything = list(set(to_add))
    all_prefixees = set(
        segments[:ix] for segments in everything for ix in range(len(segments)))
    everything_deduped = set(everything) - all_prefixees

    # Check basic addition.
    node = _BoundedTrieNode()
    total_size = node.size()
    self.assertEqual(total_size, 1)
    for segments in everything:
      total_size += node.add(segments)
    self.assertEqual(node.size(), len(everything_deduped), node)
    self.assertEqual(node.size(), total_size, node)
    self.assert_covers(node, everything_deduped)

    # Check merging
    node0 = _BoundedTrieNode()
    node0.add_all(everything[0::2])
    node1 = _BoundedTrieNode()
    node1.add_all(everything[1::2])
    pre_merge_size = node0.size()
    merge_delta = node0.merge(node1)
    self.assertEqual(node0.size(), pre_merge_size + merge_delta)
    self.assertEqual(node0, node)

    # Check trimming.
    if node.size() > 1:
      trim_delta = node.trim()
      self.assertLess(trim_delta, 0, node)
      self.assertEqual(node.size(), total_size + trim_delta)
      self.assert_covers(node, everything_deduped, max_truncated=1)

    if node.size() > 1:
      trim2_delta = node.trim()
      self.assertLess(trim2_delta, 0)
      self.assertEqual(node.size(), total_size + trim_delta + trim2_delta)
      self.assert_covers(node, everything_deduped, max_truncated=2)

    # Adding after trimming should be a no-op.
    node_copy = copy.deepcopy(node)
    for segments in everything:
      self.assertEqual(node.add(segments), 0)
    self.assertEqual(node, node_copy)

    # Merging after trimming should be a no-op.
    self.assertEqual(node.merge(node0), 0)
    self.assertEqual(node.merge(node1), 0)
    self.assertEqual(node, node_copy)

    if node._truncated:
      expected_delta = 0
    else:
      expected_delta = 2

    # Adding something new is not.
    new_values = [('new1', ), ('new2', 'new2.1')]
    self.assertEqual(node.add_all(new_values), expected_delta)
    self.assert_covers(
        node, list(everything_deduped) + new_values, max_truncated=2)

    # Nor is merging something new.
    new_values_node = _BoundedTrieNode()
    new_values_node.add_all(new_values)
    self.assertEqual(node_copy.merge(new_values_node), expected_delta)
    self.assert_covers(
        node_copy, list(everything_deduped) + new_values, max_truncated=2)

  def run_fuzz(self, iterations=10, **params):
    for _ in range(iterations):
      seed = random.getrandbits(64)
      segments = self.random_segments(**params, rand=random.Random(seed))
      try:
        self.run_test(segments)
      except:
        print("SEED", seed)
        raise

  def test_trivial(self):
    self.run_test([('a', 'b'), ('a', 'c')])

  def test_flat(self):
    self.run_test([('a', 'a'), ('b', 'b'), ('c', 'c')])

  def test_deep(self):
    self.run_test([('a', ) * 10, ('b', ) * 12])

  def test_small(self):
    self.run_fuzz(n=5, min_depth=2, max_depth=3, overlap=0.5)

  def test_medium(self):
    self.run_fuzz(n=20, min_depth=2, max_depth=4, overlap=0.5)

  def test_large_sparse(self):
    self.run_fuzz(n=120, min_depth=2, max_depth=4, overlap=0.2)

  def test_large_dense(self):
    self.run_fuzz(n=120, min_depth=2, max_depth=4, overlap=0.8)

  def test_bounded_trie_data_combine(self):
    empty = BoundedTrieData()
    # The merging here isn't complicated we're just ensuring that
    # BoundedTrieData invokes _BoundedTrieNode correctly.
    singletonA = BoundedTrieData(singleton=('a', 'a'))
    singletonB = BoundedTrieData(singleton=('b', 'b'))
    lots_root = _BoundedTrieNode()
    lots_root.add_all([('c', 'c'), ('d', 'd')])
    lots = BoundedTrieData(root=lots_root)
    self.assertEqual(empty.get_result(), set())
    self.assertEqual(
        empty.combine(singletonA).get_result(), set([('a', 'a', False)]))
    self.assertEqual(
        singletonA.combine(empty).get_result(), set([('a', 'a', False)]))
    self.assertEqual(
        singletonA.combine(singletonB).get_result(),
        set([('a', 'a', False), ('b', 'b', False)]))
    self.assertEqual(
        singletonA.combine(lots).get_result(),
        set([('a', 'a', False), ('c', 'c', False), ('d', 'd', False)]))
    self.assertEqual(
        lots.combine(singletonA).get_result(),
        set([('a', 'a', False), ('c', 'c', False), ('d', 'd', False)]))

  def test_bounded_trie_data_combine_trim(self):
    left = _BoundedTrieNode()
    left.add_all([('a', 'x'), ('b', 'd')])
    right = _BoundedTrieNode()
    right.add_all([('a', 'y'), ('c', 'd')])
    self.assertEqual(
        BoundedTrieData(root=left).combine(
            BoundedTrieData(root=right, bound=3)).get_result(),
        set([('a', True), ('b', 'd', False), ('c', 'd', False)]))

  def test_merge_on_empty_node(self):
    root1 = _BoundedTrieNode()
    root2 = _BoundedTrieNode()
    root2.add_all([["a", "b", "c"], ["a", "b", "d"], ["a", "e"]])
    self.assertEqual(2, root1.merge(root2))
    self.assertEqual(3, root1.size())
    self.assertFalse(root1._truncated)

  def test_merge_with_empty_node(self):
    root1 = _BoundedTrieNode()
    root1.add_all([["a", "b", "c"], ["a", "b", "d"], ["a", "e"]])
    root2 = _BoundedTrieNode()

    self.assertEqual(0, root1.merge(root2))
    self.assertEqual(3, root1.size())
    self.assertFalse(root1._truncated)


class TestHistogramCell(unittest.TestCase):
  @classmethod
  def _modify_histogram(cls, d):
    for i in range(cls.NUM_ITERATIONS):
      d.update(i)

  NUM_THREADS = 5
  NUM_ITERATIONS = 100

  def test_parallel_access(self):
    # We create NUM_THREADS threads that concurrently modify the distribution.
    threads = []
    bucket_type = LinearBucket(0, 1, 100)
    d = HistogramCell(bucket_type)
    for _ in range(TestHistogramCell.NUM_THREADS):
      t = threading.Thread(
          target=TestHistogramCell._modify_histogram, args=(d, ))
      threads.append(t)
      t.start()

    for t in threads:
      t.join()

    histogram = Histogram(bucket_type)
    for _ in range(self.NUM_THREADS):
      for i in range(self.NUM_ITERATIONS):
        histogram.record(i)

    self.assertEqual(d.get_cumulative(), HistogramData(histogram))

  def test_basic_operations(self):
    d = HistogramCellFactory(LinearBucket(0, 1, 10))()
    d.update(10)
    self.assertEqual(
        str(d.get_cumulative()),
        'HistogramData(Total count: 1, P99: >=10, P90: >=10, P50: >=10)')
    d.update(0)
    self.assertEqual(
        str(d.get_cumulative()),
        'HistogramData(Total count: 2, P99: >=10, P90: >=10, P50: 1)')
    d.update(5)
    self.assertEqual(
        str(d.get_cumulative()),
        'HistogramData(Total count: 3, P99: >=10, P90: >=10, P50: 6)')


if __name__ == '__main__':
  unittest.main()