# Lint as: python3

# Copyright 2020 Google LLC.
#
# Licensed 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 beam_image."""

import base64
import unittest

import apache_beam as beam
from apache_beam.testing import util
import numpy as np
import PIL
from PIL import Image

from tfrecorder import beam_image
from tfrecorder import test_utils
from tfrecorder import schema


class BeamImageTests(unittest.TestCase):
  """Tests for beam_image.py"""

  def setUp(self):
    self.pipeline = test_utils.get_test_pipeline()
    self.df = test_utils.get_test_df()
    self.image_file = 'tfrecorder/test_data/images/cat/cat-640x853-1.jpg'

  def test_load(self):
    """Tests the image loading function."""
    img = beam_image.load(self.image_file)
    self.assertIsInstance(img, PIL.JpegImagePlugin.JpegImageFile)

  def test_file_not_found_load(self):
    """Test loading an image that doesn't exist."""
    with self.assertRaises(OSError):
      _ = beam_image.load('tfrecorder/test_data/images/cat/food.jpg')

  def test_mode_to_channel(self):
    """Tests `mode_to_channel`."""

    actual = [beam_image.mode_to_channel(mode)
              for mode in ('L', 'RGB', 'whatever')]
    self.assertEqual(actual, [1, 3, 3])

  def test_channel_to_mode(self):
    """Tests `channel_to_mode`."""

    actual = [beam_image.channel_to_mode(channel) for channel in (1, 2, 3)]
    self.assertEqual(actual, ['L', 'RGB', 'RGB'])

  def test_base64_encode(self):
    """Tests encode function."""
    img = beam_image.load(self.image_file)
    enc = beam_image.encode(img)
    decode = base64.b64decode(enc, altchars=beam_image.BASE64_ALTCHARS)
    self.assertEqual(img.tobytes(), decode)

  def test_base64_decode(self):
    """Tests `decode` function."""

    image = Image.open(self.image_file)
    b64_bytes = base64.b64encode(
        image.tobytes(), altchars=beam_image.BASE64_ALTCHARS)
    width, height = image.size
    actual = beam_image.decode(b64_bytes, width, height, 3)
    np.testing.assert_array_equal(np.asarray(actual), np.asarray(image))

  def test_extract_image_dofn(self):
    """Tests ExtractImageDoFn."""

    with self.pipeline as p:

      converter = schema.get_tft_coder(['split', 'image_uri', 'label'],
                                       schema.image_csv_schema)


      extract_images_fn = beam_image.ExtractImagesDoFn('image_uri')

      data = (
          p
          | 'ReadFromDataFrame' >> beam.Create(self.df.values.tolist())
          | 'FlattenDataFrame' >> beam.Map(
              lambda x: ','.join([str(item) for item in x]))
          | 'DecodeCSV' >> beam.Map(converter.decode)
          | 'ExtractImage' >> beam.ParDo(extract_images_fn)
      )

      def key_matcher(expected_keys):
        """Custom Beam Test Matcher that tests a list of keys exist."""
        def _equal(actual):
          """ _equal raises a BeamAssertException when an element in the
              PCollection doesn't contain the image extraction keys."""
          expected_keys_ = set(expected_keys)
          for element in actual:
            actual_keys = set(element.keys())
            if actual_keys != expected_keys_:
              raise util.BeamAssertException(
                  'PCollection key match failed. Actual ({}) vs. expected ({})'
                  .format(actual_keys, expected_keys_))
        return _equal

      expected_keys = ['image_name', 'label', 'split', 'image',
                       'image_height', 'image_width', 'image_channels']
      util.assert_that(data, key_matcher(expected_keys))
