#
# 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 the stream implementations."""
# pytype: skip-file

import logging
import math
import unittest

import numpy as np

from apache_beam.coders import slow_stream


class StreamTest(unittest.TestCase):
  # pylint: disable=invalid-name
  InputStream = slow_stream.InputStream
  OutputStream = slow_stream.OutputStream
  ByteCountingOutputStream = slow_stream.ByteCountingOutputStream

  # pylint: enable=invalid-name

  def test_read_write(self):
    out_s = self.OutputStream()
    out_s.write(b'abc')
    out_s.write(b'\0\t\n')
    out_s.write(b'xyz', True)
    out_s.write(b'', True)
    in_s = self.InputStream(out_s.get())
    self.assertEqual(b'abc\0\t\n', in_s.read(6))
    self.assertEqual(b'xyz', in_s.read_all(True))
    self.assertEqual(b'', in_s.read_all(True))

  def test_read_all(self):
    out_s = self.OutputStream()
    out_s.write(b'abc')
    in_s = self.InputStream(out_s.get())
    self.assertEqual(b'abc', in_s.read_all(False))

  def test_read_write_byte(self):
    out_s = self.OutputStream()
    out_s.write_byte(1)
    out_s.write_byte(0)
    out_s.write_byte(0xFF)
    in_s = self.InputStream(out_s.get())
    self.assertEqual(1, in_s.read_byte())
    self.assertEqual(0, in_s.read_byte())
    self.assertEqual(0xFF, in_s.read_byte())

  def test_read_write_large(self):
    values = range(4 * 1024)
    out_s = self.OutputStream()
    for v in values:
      out_s.write_bigendian_int64(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_bigendian_int64())

  def run_read_write_var_int64(self, values):
    out_s = self.OutputStream()
    for v in values:
      out_s.write_var_int64(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_var_int64())

  def test_small_var_int64(self):
    self.run_read_write_var_int64(range(-10, 30))

  def test_medium_var_int64(self):
    base = -1.7
    self.run_read_write_var_int64([
        int(base**pow)
        for pow in range(1, int(63 * math.log(2) / math.log(-base)))
    ])

  def test_large_var_int64(self):
    self.run_read_write_var_int64([0, 2**63 - 1, -2**63, 2**63 - 3])

  def run_read_write_var_int32(self, values):
    out_s = self.OutputStream()
    for v in values:
      out_s.write_var_int32(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_var_int32())

  def test_small_var_int32(self):
    self.run_read_write_var_int32(range(-10, 30))

  def test_medium_var_int32(self):
    base = -1.7
    self.run_read_write_var_int32([
        int(base**pow)
        for pow in range(1, int(31 * math.log(2) / math.log(-base)))
    ])

  def test_large_var_int32(self):
    self.run_read_write_var_int32([0, 2**31 - 1, -2**31, 2**31 - 3])

  def test_read_write_double(self):
    values = 0, 1, -1, 1e100, 1.0 / 3, math.pi, float('inf')
    out_s = self.OutputStream()
    for v in values:
      out_s.write_bigendian_double(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_bigendian_double())

  def test_read_write_float(self):
    values = 0, 1, -1, 1e20, 1.0 / 3, math.pi, float('inf')
    # Restrict to single precision before coder roundtrip
    values = tuple(float(np.float32(v)) for v in values)
    out_s = self.OutputStream()
    for v in values:
      out_s.write_bigendian_float(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_bigendian_float())

  def test_read_write_bigendian_int64(self):
    values = 0, 1, -1, 2**63 - 1, -2**63, int(2**61 * math.pi)
    out_s = self.OutputStream()
    for v in values:
      out_s.write_bigendian_int64(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_bigendian_int64())

  def test_read_write_bigendian_uint64(self):
    values = 0, 1, 2**64 - 1, int(2**61 * math.pi)
    out_s = self.OutputStream()
    for v in values:
      out_s.write_bigendian_uint64(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_bigendian_uint64())

  def test_read_write_bigendian_int32(self):
    values = 0, 1, -1, 2**31 - 1, -2**31, int(2**29 * math.pi)
    out_s = self.OutputStream()
    for v in values:
      out_s.write_bigendian_int32(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_bigendian_int32())

  def test_read_write_bigendian_int16(self):
    values = 0, 1, -1, 2**15 - 1, -2**15, int(2**13 * math.pi)
    out_s = self.OutputStream()
    for v in values:
      out_s.write_bigendian_int16(v)
    in_s = self.InputStream(out_s.get())
    for v in values:
      self.assertEqual(v, in_s.read_bigendian_int16())

  def test_byte_counting(self):
    bc_s = self.ByteCountingOutputStream()
    self.assertEqual(0, bc_s.get_count())
    bc_s.write(b'def')
    self.assertEqual(3, bc_s.get_count())
    bc_s.write(b'')
    self.assertEqual(3, bc_s.get_count())
    bc_s.write_byte(10)
    self.assertEqual(4, bc_s.get_count())
    # "nested" also writes the length of the string, which should
    # cause 1 extra byte to be counted.
    bc_s.write(b'2345', nested=True)
    self.assertEqual(9, bc_s.get_count())
    bc_s.write_var_int64(63)
    self.assertEqual(10, bc_s.get_count())
    bc_s.write_bigendian_int64(42)
    self.assertEqual(18, bc_s.get_count())
    bc_s.write_bigendian_int32(36)
    self.assertEqual(22, bc_s.get_count())
    bc_s.write_bigendian_double(6.25)
    self.assertEqual(30, bc_s.get_count())
    bc_s.write_bigendian_uint64(47)
    self.assertEqual(38, bc_s.get_count())


try:
  # pylint: disable=wrong-import-position
  from apache_beam.coders import stream

  class FastStreamTest(StreamTest):
    """Runs the test with the compiled stream classes."""
    InputStream = stream.InputStream
    OutputStream = stream.OutputStream
    ByteCountingOutputStream = stream.ByteCountingOutputStream

  class SlowFastStreamTest(StreamTest):
    """Runs the test with compiled and uncompiled stream classes."""
    InputStream = stream.InputStream
    OutputStream = slow_stream.OutputStream
    ByteCountingOutputStream = slow_stream.ByteCountingOutputStream

  class FastSlowStreamTest(StreamTest):
    """Runs the test with uncompiled and compiled stream classes."""
    InputStream = slow_stream.InputStream
    OutputStream = stream.OutputStream
    ByteCountingOutputStream = stream.ByteCountingOutputStream

except ImportError:
  pass

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