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

"""Tests for apache_beam.runners.interactive.pipeline_ananlyzer.

This module is experimental. No backwards-compatibility guarantees.
"""

# pytype: skip-file

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import unittest

import apache_beam as beam
from apache_beam.runners.direct import direct_runner
from apache_beam.runners.interactive import cache_manager as cache
from apache_beam.runners.interactive import pipeline_analyzer


def to_stable_runner_api(p):
  """The extra round trip ensures a stable pipeline proto.
  """
  return (
      beam.pipeline.Pipeline.from_runner_api(
          p.to_runner_api(use_fake_coders=True), p.runner,
          p._options).to_runner_api(use_fake_coders=True))


class PipelineAnalyzerTest(unittest.TestCase):
  def setUp(self):
    self.runner = direct_runner.DirectRunner()
    self.cache_manager = cache.FileBasedCacheManager()

  def tearDown(self):
    self.cache_manager.cleanup()

  def assertPipelineEqual(self, pipeline_proto1, pipeline_proto2):
    """A naive check for Pipeline proto equality.
    """
    components1 = pipeline_proto1.components
    components2 = pipeline_proto2.components

    self.assertEqual(len(components1.transforms), len(components2.transforms))
    self.assertEqual(
        len(components1.pcollections), len(components2.pcollections))

    # GreatEqual instead of Equal because the pipeline_proto_to_execute could
    # include more windowing_stratagies and coders than necessary.
    self.assertGreaterEqual(
        len(components1.windowing_strategies),
        len(components2.windowing_strategies))
    self.assertGreaterEqual(len(components1.coders), len(components2.coders))
    self.assertTransformEqual(
        pipeline_proto1,
        pipeline_proto1.root_transform_ids[0],
        pipeline_proto2,
        pipeline_proto2.root_transform_ids[0])

  def assertTransformEqual(
      self, pipeline_proto1, transform_id1, pipeline_proto2, transform_id2):
    """A naive check for Transform proto equality.
    """
    transform_proto1 = pipeline_proto1.components.transforms[transform_id1]
    transform_proto2 = pipeline_proto2.components.transforms[transform_id2]
    self.assertEqual(transform_proto1.spec.urn, transform_proto2.spec.urn)
    # Skipping payload checking because PTransforms of the same functionality
    # could generate different payloads.
    self.assertEqual(
        len(transform_proto1.subtransforms),
        len(transform_proto2.subtransforms))
    self.assertSetEqual(
        set(transform_proto1.inputs), set(transform_proto2.inputs))
    self.assertSetEqual(
        set(transform_proto1.outputs), set(transform_proto2.outputs))

  def test_basic(self):
    p = beam.Pipeline(runner=self.runner)

    # The cold run.
    pcoll = (
        p
        | 'Create' >> beam.Create([1, 2, 3])
        | 'Double' >> beam.Map(lambda x: x * 2)
        | 'Square' >> beam.Map(lambda x: x**2))
    analyzer = pipeline_analyzer.PipelineAnalyzer(
        self.cache_manager, to_stable_runner_api(p), self.runner)
    pipeline_to_execute = beam.pipeline.Pipeline.from_runner_api(
        analyzer.pipeline_proto_to_execute(), p.runner, p._options)
    pipeline_to_execute.run().wait_until_finish()
    self.assertEqual(
        len(analyzer.tl_required_trans_ids()),
        7  # Create, Double, Square, CacheSample * 3, CacheFull
    )
    self.assertEqual(len(analyzer.tl_referenced_pcoll_ids()), 3)
    self.assertEqual(len(analyzer.read_cache_ids()), 0)
    self.assertEqual(len(analyzer.write_cache_ids()), 4)

    # The second run.
    _ = (
        pcoll
        | 'Triple' >> beam.Map(lambda x: x * 3)
        | 'Cube' >> beam.Map(lambda x: x**3))
    analyzer = pipeline_analyzer.PipelineAnalyzer(
        self.cache_manager, to_stable_runner_api(p), self.runner)
    pipeline_to_execute = beam.pipeline.Pipeline.from_runner_api(
        analyzer.pipeline_proto_to_execute(), p.runner, p._options)
    self.assertEqual(
        len(analyzer.tl_required_trans_ids()),
        6  # Read, Triple, Cube, CacheSample * 2, CacheFull
    )
    self.assertEqual(len(analyzer.tl_referenced_pcoll_ids()), 3)
    self.assertEqual(len(analyzer.read_cache_ids()), 1)
    self.assertEqual(len(analyzer.write_cache_ids()), 3)

    # No need to actually execute the second run.

  def test_word_count(self):
    p = beam.Pipeline(runner=self.runner)

    class WordExtractingDoFn(beam.DoFn):
      def process(self, element):
        text_line = element.strip()
        words = text_line.split()
        return words

    # Count the occurrences of each word.
    pcoll1 = p | beam.Create(['to be or not to be that is the question'])
    pcoll2 = pcoll1 | 'Split' >> beam.ParDo(WordExtractingDoFn())
    pcoll3 = pcoll2 | 'Pair with One' >> beam.Map(lambda x: (x, 1))
    pcoll4 = pcoll3 | 'Group' >> beam.GroupByKey()
    pcoll5 = pcoll4 | 'Count' >> beam.Map(lambda item: (item[0], sum(item[1])))

    analyzer = pipeline_analyzer.PipelineAnalyzer(
        self.cache_manager, to_stable_runner_api(p), self.runner)

    cache_label1 = 'PColl-1111111'
    cache_label2 = 'PColl-2222222'
    cache_label3 = 'PColl-3333333'
    cache_label4 = 'PColl-4444444'
    cache_label5 = 'PColl-5555555'

    # pylint: disable=expression-not-assigned
    pcoll1 | 'CacheSample%s' % cache_label1 >> cache.WriteCache(
        self.cache_manager, cache_label1, sample=True, sample_size=10)
    pcoll2 | 'CacheSample%s' % cache_label2 >> cache.WriteCache(
        self.cache_manager, cache_label2, sample=True, sample_size=10)
    pcoll3 | 'CacheSample%s' % cache_label3 >> cache.WriteCache(
        self.cache_manager, cache_label3, sample=True, sample_size=10)
    pcoll4 | 'CacheSample%s' % cache_label4 >> cache.WriteCache(
        self.cache_manager, cache_label3, sample=True, sample_size=10)
    pcoll5 | 'CacheSample%s' % cache_label5 >> cache.WriteCache(
        self.cache_manager, cache_label3, sample=True, sample_size=10)
    pcoll5 | 'CacheFull%s' % cache_label5 >> cache.WriteCache(
        self.cache_manager, cache_label3)
    expected_pipeline_proto = to_stable_runner_api(p)

    self.assertPipelineEqual(
        analyzer.pipeline_proto_to_execute(), expected_pipeline_proto)

    pipeline_to_execute = beam.pipeline.Pipeline.from_runner_api(
        analyzer.pipeline_proto_to_execute(), p.runner, p._options)
    pipeline_to_execute.run().wait_until_finish()

  def test_write_cache_expansion(self):
    p = beam.Pipeline(runner=self.runner)

    pcoll1 = p | 'Create' >> beam.Create([1, 2, 3])
    pcoll2 = pcoll1 | 'Double' >> beam.Map(lambda x: x * 2)
    pcoll3 = pcoll2 | 'Square' >> beam.Map(lambda x: x**2)
    analyzer = pipeline_analyzer.PipelineAnalyzer(
        self.cache_manager, to_stable_runner_api(p), self.runner)

    cache_label1 = 'PColl-1234567'
    cache_label2 = 'PColl-7654321'
    cache_label3 = 'PColl-3141593'

    # pylint: disable=expression-not-assigned
    pcoll1 | 'CacheSample%s' % cache_label1 >> cache.WriteCache(
        self.cache_manager, cache_label1, sample=True, sample_size=10)
    pcoll2 | 'CacheSample%s' % cache_label2 >> cache.WriteCache(
        self.cache_manager, cache_label2, sample=True, sample_size=10)
    pcoll3 | 'CacheSample%s' % cache_label3 >> cache.WriteCache(
        self.cache_manager, cache_label3, sample=True, sample_size=10)
    pcoll3 | 'CacheFull%s' % cache_label3 >> cache.WriteCache(
        self.cache_manager, cache_label3)
    expected_pipeline_proto = to_stable_runner_api(p)

    # Since WriteCache expansion leads to more than 50 PTransform protos in the
    # pipeline, a simple check of proto map size is enough.
    self.assertPipelineEqual(
        analyzer.pipeline_proto_to_execute(), expected_pipeline_proto)

  def test_read_cache_expansion(self):
    p = beam.Pipeline(runner=self.runner)

    # The cold run.
    pcoll = (
        p
        | 'Create' >> beam.Create([1, 2, 3])
        | 'Double' >> beam.Map(lambda x: x * 2)
        | 'Square' >> beam.Map(lambda x: x**2))
    pipeline_proto = to_stable_runner_api(p)

    pipeline_info = pipeline_analyzer.PipelineInfo(pipeline_proto.components)
    pcoll_id = 'ref_PCollection_PCollection_10'  # Output PCollection of Square
    cache_label1 = pipeline_info.cache_label(pcoll_id)

    analyzer = pipeline_analyzer.PipelineAnalyzer(
        self.cache_manager, pipeline_proto, self.runner)
    pipeline_to_execute = beam.pipeline.Pipeline.from_runner_api(
        analyzer.pipeline_proto_to_execute(), p.runner, p._options)
    pipeline_to_execute.run().wait_until_finish()

    # The second run.
    _ = (
        pcoll
        | 'Triple' >> beam.Map(lambda x: x * 3)
        | 'Cube' >> beam.Map(lambda x: x**3))
    analyzer = pipeline_analyzer.PipelineAnalyzer(
        self.cache_manager, to_stable_runner_api(p), self.runner)

    expected_pipeline = beam.Pipeline(runner=self.runner)
    pcoll1 = (
        expected_pipeline
        | 'Load%s' % cache_label1 >> cache.ReadCache(
            self.cache_manager, cache_label1))
    pcoll2 = pcoll1 | 'Triple' >> beam.Map(lambda x: x * 3)
    pcoll3 = pcoll2 | 'Cube' >> beam.Map(lambda x: x**3)

    cache_label2 = 'PColl-7654321'
    cache_label3 = 'PColl-3141593'

    # pylint: disable=expression-not-assigned
    pcoll2 | 'CacheSample%s' % cache_label2 >> cache.WriteCache(
        self.cache_manager, cache_label2, sample=True, sample_size=10)
    pcoll3 | 'CacheSample%s' % cache_label3 >> cache.WriteCache(
        self.cache_manager, cache_label3, sample=True, sample_size=10)
    pcoll3 | 'CacheFull%s' % cache_label3 >> cache.WriteCache(
        self.cache_manager, cache_label3)

    # Since ReadCache & WriteCache expansion leads to more than 50 PTransform
    # protos in the pipeline, a simple check of proto map size is enough.
    self.assertPipelineEqual(
        analyzer.pipeline_proto_to_execute(),
        to_stable_runner_api(expected_pipeline))


class PipelineInfoTest(unittest.TestCase):
  def setUp(self):
    self.runner = direct_runner.DirectRunner()

  def test_passthrough(self):
    """
    Test that PTransforms which pass through their input PCollection can be
    used with PipelineInfo.
    """
    class Passthrough(beam.PTransform):
      def expand(self, pcoll):
        return pcoll

    p = beam.Pipeline(runner=self.runner)
    p | beam.Impulse() | Passthrough()  # pylint: disable=expression-not-assigned
    proto = to_stable_runner_api(p).components
    info = pipeline_analyzer.PipelineInfo(proto)
    for pcoll_id in info.all_pcollections():
      # FIXME: If PipelineInfo does not support passthrough PTransforms, this
      #        will only fail some of the time, depending on the ordering of
      #        transforms in the Pipeline proto.

      # Should not throw exception
      info.cache_label(pcoll_id)


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