#
# 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 batched type converters."""

import contextlib
import random
import typing
import unittest

import numpy as np
from parameterized import parameterized
from parameterized import parameterized_class

from apache_beam.typehints import typehints
from apache_beam.typehints.batch import BatchConverter
from apache_beam.typehints.batch import N
from apache_beam.typehints.batch import NumpyArray


@parameterized_class([
    {
        'batch_typehint': np.ndarray,
        'element_typehint': np.int32,
        'batch': np.array(range(100), np.int32)
    },
    {
        'batch_typehint': NumpyArray[np.int64, (N, 10)],
        'element_typehint': NumpyArray[np.int64, (10, )],
        'batch': np.array([list(range(i, i + 10)) for i in range(100)],
                          np.int64),
    },
    {
        'batch_typehint': typehints.List[str],
        'element_typehint': str,
        'batch': ["foo" * (i % 5) + str(i) for i in range(1000)],
    },
    {
        'batch_typehint': typing.List[str],
        'element_typehint': str,
        'batch': ["foo" * (i % 5) + str(i) for i in range(1000)],
    },
])
class BatchConverterTest(unittest.TestCase):
  def create_batch_converter(self):
    return BatchConverter.from_typehints(
        element_type=self.element_typehint, batch_type=self.batch_typehint)

  def setUp(self):
    self.converter = self.create_batch_converter()
    self.normalized_batch_typehint = typehints.normalize(self.batch_typehint)
    self.normalized_element_typehint = typehints.normalize(
        self.element_typehint)

  def equality_check(self, left, right):
    if isinstance(left, np.ndarray) and isinstance(right, np.ndarray):
      return np.array_equal(left, right)
    else:
      return left == right

  def test_typehint_validates(self):
    typehints.validate_composite_type_param(self.batch_typehint, '')
    typehints.validate_composite_type_param(self.element_typehint, '')

  def test_type_check(self):
    typehints.check_constraint(self.normalized_batch_typehint, self.batch)

  def test_type_check_element(self):
    for element in self.converter.explode_batch(self.batch):
      typehints.check_constraint(self.normalized_element_typehint, element)

  def test_explode_rebatch(self):
    exploded = list(self.converter.explode_batch(self.batch))
    rebatched = self.converter.produce_batch(exploded)

    typehints.check_constraint(self.normalized_batch_typehint, rebatched)
    self.assertTrue(self.equality_check(self.batch, rebatched))

  def test_estimate_byte_size_implemented(self):
    # Just verify that we can call byte size
    self.assertGreater(self.converter.estimate_byte_size(self.batch), 0)

  @parameterized.expand([
      (2, ),
      (3, ),
      (10, ),
  ])
  def test_estimate_byte_size_partitions(self, N):
    elements = list(self.converter.explode_batch(self.batch))

    # Split elements into N contiguous partitions, create a batch out of each
    batches = [
        self.converter.produce_batch(
            elements[len(elements) * i // N:len(elements) * (i + 1) // N])
        for i in range(N)
    ]

    # Some estimate_byte_size implementations use random samples,
    # set a seed temporarily to make this test deterministic
    with temp_seed(12345):
      partitioned_size_estimate = sum(
          self.converter.estimate_byte_size(batch) for batch in batches)
      size_estimate = self.converter.estimate_byte_size(self.batch)

    # Assert that size estimate for partitions is within 10% of size estimate
    # for the whole partition.
    self.assertLessEqual(
        abs(partitioned_size_estimate / size_estimate - 1), 0.1)

  @parameterized.expand([
      (2, ),
      (3, ),
      (10, ),
  ])
  def test_combine_batches(self, N):
    elements = list(self.converter.explode_batch(self.batch))

    # Split elements into N contiguous partitions, create a batch out of each
    batches = [
        self.converter.produce_batch(
            elements[len(elements) * i // N:len(elements) * (i + 1) // N])
        for i in range(N)
    ]

    # Combine the batches, output should be equivalent to the original batch
    combined = self.converter.combine_batches(batches)

    self.assertTrue(self.equality_check(self.batch, combined))

  def test_equals(self):
    self.assertTrue(self.converter == self.create_batch_converter())
    self.assertTrue(self.create_batch_converter() == self.converter)

  def test_hash(self):
    self.assertEqual(hash(self.create_batch_converter()), hash(self.converter))


class BatchConverterErrorsTest(unittest.TestCase):
  @parameterized.expand([
      (
          typing.List[int],
          str,
          r'batch type must be List\[T\] for element type T',
      ),
      (
          np.ndarray,
          typing.Any,
          r'Element type is not a dtype',
      ),
      (
          np.array,
          np.int64,
          (
              r'batch type must be np\.ndarray or '
              r'beam\.typehints\.batch\.NumpyArray\[\.\.\]'),
      ),
      (
          NumpyArray[np.int64, (3, N, 2)],
          NumpyArray[np.int64, (3, 7)],
          r'Failed to align batch type\'s batch dimension',
      ),
  ])
  def test_construction_errors(
      self, batch_typehint, element_typehint, error_regex):
    with self.assertRaisesRegex(TypeError, error_regex):
      BatchConverter.from_typehints(
          element_type=element_typehint, batch_type=batch_typehint)


@contextlib.contextmanager
def temp_seed(seed):
  state = random.getstate()
  random.seed(seed)
  try:
    yield
  finally:
    random.setstate(state)


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